更好的方法通过键合并字典列表

Question

我有一个字典列表，还有一个函数可以从这些字典中提取值。我的目标是得到一个新的字典，这个字典的键是我用这个函数处理原字典列表后得到的值。而这个新字典的值则是那些从原字典列表中提取出来的字典子集，这些字典的键正好对应函数返回的值。

我知道这个解释听起来很复杂，所以我用一个实现来展示一下：

keygen = lambda x: x['key']

data = [{'key': 'key1',
         'data': 'value2'},
        {'key': 'key3',
         'data': 'value2'},
        {'key': 'key2',
         'data': 'value2'},
        {'key': 'key2',
         'data': 'value2'},
        {'key': 'key1',
         'data': 'value2'}]

def merge_by_keygen(data, keygen):
    return_value = {} 
    for dataset in data:
        if keygen(dataset) not in return_value.keys():
            return_value[keygen(dataset)] = [] 
        return_value[keygen(dataset)].append(dataset)
    return return_value

merge_by_keygen(data, keygen)

返回结果：

{'key3': [{'data': 'value2', 'key': 'key3'}], 
 'key2': [{'data': 'value2', 'key': 'key2'}, {'data': 'value2', 'key': 'key2'}], 
 'key1': [{'data': 'value2', 'key': 'key1'}, {'data': 'value2', 'key': 'key1'}]}

我希望能找到一种更简洁、更紧凑的实现方式，比如用字典或列表推导式。谢谢！

Answer 1

我觉得这样可以实现

return_value = {}
for d in data:
    return_value.setdefault(keygen(d), []).append(d)

你可以用列表推导式来写，但用列表推导式的副作用来影响数据，然后再生成一堆None的结果然后丢掉，这样写起来就不好看...

r = {}
[r.setdefault(keygen(d), []).append(d) for d in data]

你函数的核心其实就是用到了字典的setdefault方法。前面提到的三行代码，其实就是在调用keygen，检查这个键是否在返回的字典里，如果没有，就创建一个空列表，把这个空列表存到字典里，然后再查询字典，准备把列表添加进去——这些操作都是通过setdefault()来完成的。

Answer 2

我觉得这个问题不太好理解，不过你可以用一个叫 collections.defaultdict(list) 的东西来让代码看起来更整洁：

import collections

def merge_by_keygen(data, keygen):
    return_value = collections.defaultdict(list)
    for dataset in data:
        key = keygen(dataset)
        return_value[key].append(dataset)
    return return_value

对我来说，这样看起来挺干净的。如果你想的话，可以试着调整一下调用 keygen 函数的位置，但我觉得这样可能会让代码变得不那么清晰。

Answer 3

如果你不介意使用一个第三方的工具包，这个操作可以很简单地通过 toolz.groupby 来完成：

>>> import toolz
>>> toolz.groupby(keygen, data)
{'key1': [{'data': 'value2', 'key': 'key1'},
          {'data': 'value2', 'key': 'key1'}],
 'key2': [{'data': 'value2', 'key': 'key2'},
          {'data': 'value2', 'key': 'key2'}],
 'key3': [{'data': 'value2', 'key': 'key3'}]}

使用 toolz.groupby('key', data) 也能得到相同的结果。

Answer 4

这是一个很适合用 itertools.groupby 来解决的问题。

实现方法

from itertools import groupby
from operator import itemgetter
groups = groupby(sorted(data, key = itemgetter('key')), key = itemgetter('key'))
data_dict = {k : list(g) for k, g in groups}

或者如果你喜欢一行代码的写法

data_dict = {k : list(g) 
             for k, g in groupby(sorted(data, 
                                        key = itemgetter('key')), 
                                 key = itemgetter('key'))}

输出结果

{'key1': [{'data': 'value2', 'key': 'key1'},
          {'data': 'value2', 'key': 'key1'}],
 'key2': [{'data': 'value2', 'key': 'key2'},
          {'data': 'value2', 'key': 'key2'}],
 'key3': [{'data': 'value2', 'key': 'key3'}]}