Get a dictionary of lists of values from one list based on common values from another list

user_14197646 0
Viewed 61

I have two lists. One gives me the IDs, say:

list1 = [0, 0, 0, 0, 1, 1, 2, 3, 3, 3]

Another gives me values associated with those IDs, say:

list2 = [10, 20, 30, 40, 1, 5, 9, 10, 15, 20]

I want to organize the values of list2 based on the values in list1, such that:

dict = {0:[10,20,30,40], 1:[1, 5], 2:[9], 3:[10,15,20]}

My current problem is that I need an efficient way to do this since my lists are ~150000 points long and my solution is very slow:

dict = {i:list2[np.where(np.array(obid)==i)] for i in np.unique(list1)}

Seems like a simple problem and perhaps there is an obvious answer but I just can't seem to find the right combination of words to get a good answer.

numpy python
3 Answers

If list1 is sorted (or if you pre-sort) you can use groupby to collect them into subgroups on the fly, then create a dictionary where the key is from list1 and value is from the group of list2

>>> from itertools import groupby
>>> {k:[i[1] for i in g] for k,g in groupby(zip(list1, list2), key=lambda i: i[0])}
{0: [10, 20, 30, 40], 1: [1, 5], 2: [9], 3: [10, 15, 20]}
user_22964580

Are you open to Pandas:

import pandas as pd
out_dict = {k:list(v) for k,v in pd.Series(list2).groupby(list1)}

Output:

out_dict = {k:list(v) for k,v in pd.Series(list2).groupby(list1)}

Update: for a vanilla Python:

def python_func(list1,list2):
    out_dict = {}
    for k,v in zip(list1, list2):
        if k in out_dict:
            out_dict[k].append(v)
        else:
            out_dict[k] = [v]

    return out_dict

Performance Test:

Data:

n=100000
list1 = np.random.randint(0,10,n)
list2 = np.random.randint(0,60, n)

Run time:

%timeit -n 10 out_dict = {k:list(v) for k,v in pd.Series(list2).groupby(list1)}
# 8.51 ms ± 95.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit -n 10 python_func(list1, list2);
# slightly slower, but equivalent to Stephan's answer
# 31.6 ms ± 777 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit -n 10 {k:[i[1] for i in g] for k,g in groupby(zip(list1, list2), key=lambda i: i[0])}
# 58.2 ms ± 1.31 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
user_42384080

In a single loop you can do the following using defaultdicts with no assumptions on the ordering of the indices-array:

In [1]: import numpy as np; from collections import defaultdict

In [2]: indices = np.asarray([0, 0, 0, 0, 1, 1, 2, 3, 3, 3])

In [3]: values = np.asarray([10, 20, 30, 40, 1, 5, 9, 10, 15, 20])

In [4]: mapping = defaultdict(list)
   ...: for i, v in zip(indices, values):
   ...:     mapping[i].append(v)
   ...:

In [5]: mapping
Out[5]: defaultdict(list, {0: [10, 20, 30, 40], 1: [1, 5], 2: [9], 3: [10, 15, 20]})
user_129982050
Related