CODEWITHSUNDEEP
pythonpython-2.7
sashkello
sashkello

Reputation: 17871

Splitting dict by value of one of the keys

I've got a dictionary with data of the same length (but different types), something like:

data = {
    "id": [1,1,2,2,1,2,1,2], 
    "info": ["info1","info2","info3","info4","info5","info6","info7","info8"],       
    "number": [1,2,3,4,5,6,7,8]
}

Now I'd like to split it in two by id, keeping the respective info and number. That is, to have two dicts data1 and data2.

Note: this is merely a sample, there are multiple keys in the dict and I would want to avoid using the key names, but rather loop through all of them.

What is a Pythonic way of doing it?

Upvotes: 4

Views: 8730

Answers (4)

ndpu
ndpu

Reputation: 22561

from collections import defaultdict

ids = data.pop('id')
databyid = defaultdict(lambda: defaultdict(list))

for id, values in zip(ids, zip(*data.values())):
    for kid, kval in enumerate(data.keys()):
        databyid[id][kval].append(values[kid])

if you need data in original state (with id):

 data['id'] = ids

result:

>>> databyid[1]
defaultdict(<type 'list'>, {'info': ['info1', 'info2', 'info5', 'info7'], 'number': [1, 2, 5, 7]})
>>> databyid[2]
defaultdict(<type 'list'>, {'info': ['info3', 'info4', 'info6', 'info8'], 'number': [3, 4, 6, 8]})
>>>

Upvotes: 0

shx2
shx2

Reputation: 64318

For working with records, I personally like numpy.recarray.

In [3]: import numpy as np
In [4]: fields = data.keys()
In [8]: recs = zip(*[ lst for k, lst in data.iteritems() ])

In [9]: recs[0]
Out[9]: ('info1', 1, 1)
In [10]: recs[1]
Out[10]: ('info2', 1, 2)

In [21]: ra = np.rec.fromrecords(recs, names = fields )
In [17]: ra
rec.array([('info1', 1, 1), ('info2', 1, 2), ('info3', 2, 3), ('info4', 2, 4),
       ('info5', 1, 5), ('info6', 2, 6), ('info7', 1, 7), ('info8', 2, 8)], 
      dtype=[('info', 'S5'), ('id', '<i8'), ('number', '<i8')])

In [23]: ra[ra.id == 2]
rec.array([('info3', 2, 3), ('info4', 2, 4), ('info6', 2, 6), ('info8', 2, 8)], 
      dtype=[('info', 'S5'), ('id', '<i8'), ('number', '<i8')])

In [24]: ra[ra.id == 2].number
Out[24]: array([3, 4, 6, 8])

In [25]: ra[ra.id == 2][0]
Out[25]: ('info3', 2, 3)

In [26]: ra[ra.id == 2][0].number
Out[26]: 3

If you want to group the records by id in a dict, do:

{ id: ra[ra.id == id] for id in set(ra.id) }

Upvotes: 2

lucasg
lucasg

Reputation: 11002

with comprehension lists : 

data1 = [ data["info"][idx] for idx, x in enumerate(data["id"]) if x == 1 ]
#data1 = ['info1', 'info2', 'info5', 'info7']

If you want to recover all the keys : 

data1 = [ { key : data[key][idx] for key in data.keys() }  for idx, x in enu
merate(data["id"]) if x == 1 ]
>>> data1
[{'info': 'info1', 'id': 1, 'number': 1}, {'info': 'info2', 'id': 1, 'number': 2
}, {'info': 'info5', 'id': 1, 'number': 5}, {'info': 'info7', 'id': 1, 'number':
 7}]

Upvotes: 3

AI Generated Response
AI Generated Response

Reputation: 8835

>>> from collections import defaultdict
>>> res = defaultdict(list)
>>> for ID,info in zip(data["id"],data["info"]):
    res[ID].append(info)


>>> res
defaultdict(<type 'list'>, {1: ['info1', 'info2', 'info5', 'info7'], 2: ['info3', 'info4', 'info6', 'info8']})
>>>

Upvotes: 0

Related Questions