Loop through pandas columns and append a dict of sets?

Viewed 1358

I am looking to loop through around 10 million rows in a pandas dataframe and add them to an already existing dict of sets.

For example for a dict like this

x = {10: {1, 2, 3, 5}, 12: {6, 7, 8, 9, 10}}

And a dataframe like this:

d = {'ID': [10, 10, 10, 12, 12, 12], 'Another_ID': [1, 4, 6, 6, 7, 13]}
df = pd.DataFrame(data=d)

ID   Another_ID

10   1
10   4
10   6
12   6
12   7
12   13

I would like to go through the rows and add the new values that ID "hasnt seen yet." I would like a result like this.

x = {10: {1, 2, 3, 4, 5, 6}, 12: {6, 7, 8, 9, 10, 13}}

I have tried iterating through using a simple function like the following.

for i in df [['ID' , 'Another_ID' ]] .values():
    dict[i[0]].add(i[1])

I can manually add in the values by saying the following like this, but cant do it in a loop!

  dict[10].add(6)

If anyone knows how to loop through these two pandas columns and add new values to the set, please let me know!

  • keep in mind this must be done relatively quickly as there are 10 million rows

Thanks!

3 Answers

You can convert df into a similar format as "x" using groupby and agg:

x2 = df.groupby('ID')['Another_ID'].agg(set).to_dict()
print (x2)
# {10: {1, 4, 6}, 12: {6, 7, 13}}

Now, we merge the two dictionaries using a single expression:

x3 = {k: x.get(k, set()) | x2.get(k, set()) for k in x}
print (x3)
# {10: {1, 2, 3, 4, 5, 6}, 12: {6, 7, 8, 9, 10, 13}}

Or, for an in-place merge (makes more sense if x is large and x2 is small):

for k in x2:
    x[k] = x2[k] | x.get(k, set())

print (x)
# {10: {1, 2, 3, 4, 5, 6}, 12: {6, 7, 8, 9, 10, 13}}

Where the | operator represents set union of two set operands.

One way by panda explode

out = pd.Series(x).map(list).explode().append(df.set_index('ID')['Another_ID']).groupby(level=0).agg(set).to_dict()
Out[361]: {10: {1, 2, 3, 4, 5, 6}, 12: {6, 7, 8, 9, 10, 13}}

You could treat your dataframe as a dictionary, use defaultdict to get your data out of the Pandas dataframe, and then iterate throught the dictionary to get your final output:

from collections import defaultdict

dd = defaultdict(list)

for ID, another_ID in zip(df.ID, df.Another_ID):
    dd[ID].append(another_ID)

dd

defaultdict(list, {10: [1, 4, 6], 12: [6, 7, 13]})

Final result:

{key: value.union(dd[key]) for key, value in x.items()}

{10: {1, 2, 3, 4, 5, 6}, 12: {6, 7, 8, 9, 10, 13}}
Related