Python 分组

Question

假设我有一组数据对，其中索引0是值，索引1是类型：

input = [
          ('11013331', 'KAT'), 
          ('9085267',  'NOT'), 
          ('5238761',  'ETH'), 
          ('5349618',  'ETH'), 
          ('11788544', 'NOT'), 
          ('962142',   'ETH'), 
          ('7795297',  'ETH'), 
          ('7341464',  'ETH'), 
          ('9843236',  'KAT'), 
          ('5594916',  'ETH'), 
          ('1550003',  'ETH')
        ]

我想根据它们的类型（也就是第一个索引的字符串）把它们分组，像这样：

result = [ 
           { 
             'type': 'KAT', 
             'items': ['11013331', '9843236'] 
           },
           {
             'type': 'NOT', 
             'items': ['9085267', '11788544'] 
           },
           {
             'type': 'ETH', 
             'items': ['5238761', '962142', '7795297', '7341464', '5594916', '1550003'] 
           }
         ] 

我该如何高效地做到这一点呢？

Answer 1

这个回答和@PaulMcG的回答类似，但不需要对输入进行排序。

如果你对函数式编程感兴趣，可以用一行代码（不包括导入的部分）来写出groupBy，而且和itertools.groupby不同，它不需要输入数据是排好序的：

from functools import reduce # import needed for python3; builtin in python2
from collections import defaultdict

def groupBy(key, seq):
 return reduce(lambda grp, val: grp[key(val)].append(val) or grp, seq, defaultdict(list))

在lambda中出现的... or grp是因为为了让这个reduce()正常工作，lambda需要返回它的第一个参数；因为list.append()总是返回None，所以or会一直返回grp。也就是说，这是一个小技巧，用来绕过Python的限制，确保lambda只能计算一个表达式。

这个函数会返回一个字典，字典的键是通过给定的函数计算得出的，值是原始项目的列表，顺序和原来一样。对于提问者的例子，调用这个函数groupBy(lambda pair: pair[1], input)将返回这个字典：

{'KAT': [('11013331', 'KAT'), ('9843236', 'KAT')],
 'NOT': [('9085267', 'NOT'), ('11788544', 'NOT')],
 'ETH': [('5238761', 'ETH'), ('5349618', 'ETH'), ('962142', 'ETH'), ('7795297', 'ETH'), ('7341464', 'ETH'), ('5594916', 'ETH'), ('1550003', 'ETH')]}

根据@PaulMcG的回答，提问者想要的格式可以通过将其放在列表推导式中来实现。所以这样做就可以了：

result = {key: [pair[0] for pair in values],
          for key, values in groupBy(lambda pair: pair[1], input).items()}

Answer 2

Python自带的 itertools 模块里有一个 groupby 函数，不过在使用这个函数之前，想要分组的元素必须先排好序，这样才能保证要分组的元素在列表中是连续的：

from operator import itemgetter
sortkeyfn = itemgetter(1)
input = [('11013331', 'KAT'), ('9085267', 'NOT'), ('5238761', 'ETH'), 
 ('5349618', 'ETH'), ('11788544', 'NOT'), ('962142', 'ETH'), ('7795297', 'ETH'), 
 ('7341464', 'ETH'), ('9843236', 'KAT'), ('5594916', 'ETH'), ('1550003', 'ETH')] 
input.sort(key=sortkeyfn)

现在输入看起来是这样的：

[('5238761', 'ETH'), ('5349618', 'ETH'), ('962142', 'ETH'), ('7795297', 'ETH'),
 ('7341464', 'ETH'), ('5594916', 'ETH'), ('1550003', 'ETH'), ('11013331', 'KAT'),
 ('9843236', 'KAT'), ('9085267', 'NOT'), ('11788544', 'NOT')]

groupby 会返回一系列的二元组，格式是 (key, values_iterator)。我们想要的是把这个变成一个字典的列表，其中 'type' 是键，'items' 是从 values_iterator 返回的元组中的第一个元素组成的列表。像这样：

from itertools import groupby
result = []
for key,valuesiter in groupby(input, key=sortkeyfn):
    result.append(dict(type=key, items=list(v[0] for v in valuesiter)))

现在 result 里包含了你想要的字典，正如你问题中所说的那样。

不过你可以考虑，把这些合并成一个字典，以 'type' 为键，每个值包含一个值的列表。这样的话，在你现在的形式中，要找到某个特定类型的值，你需要遍历列表，找到包含匹配 'type' 键的字典，然后再从中获取 'items' 元素。如果你用一个字典而不是一个包含单个字典的列表，你就可以通过一次查找直接找到特定类型的项目。使用 groupby 的话，这样写：

result = {}
for key,valuesiter in groupby(input, key=sortkeyfn):
    result[key] = list(v[0] for v in valuesiter)

现在 result 包含了这个字典（这和 @KennyTM 的回答中的中间 res defaultdict 类似）：

{'NOT': ['9085267', '11788544'], 
 'ETH': ['5238761', '5349618', '962142', '7795297', '7341464', '5594916', '1550003'], 
 'KAT': ['11013331', '9843236']}

（如果你想把这个简化成一行代码，你可以：

result = dict((key,list(v[0] for v in valuesiter)
              for key,valuesiter in groupby(input, key=sortkeyfn))

或者使用新式的字典推导形式：

result = {key:list(v[0] for v in valuesiter)
              for key,valuesiter in groupby(input, key=sortkeyfn)}

Answer 3

分两步来做。首先，创建一个字典。

>>> input = [('11013331', 'KAT'), ('9085267', 'NOT'), ('5238761', 'ETH'), ('5349618', 'ETH'), ('11788544', 'NOT'), ('962142', 'ETH'), ('7795297', 'ETH'), ('7341464', 'ETH'), ('9843236', 'KAT'), ('5594916', 'ETH'), ('1550003', 'ETH')]
>>> from collections import defaultdict
>>> res = defaultdict(list)
>>> for v, k in input: res[k].append(v)
...

然后，把这个字典转换成你想要的格式。

>>> [{'type':k, 'items':v} for k,v in res.items()]
[{'items': ['9085267', '11788544'], 'type': 'NOT'}, {'items': ['5238761', '5349618', '962142', '7795297', '7341464', '5594916', '1550003'], 'type': 'ETH'}, {'items': ['11013331', '9843236'], 'type': 'KAT'}]

---

也可以用itertools.groupby来实现，但这要求输入的数据必须先排好序。

>>> sorted_input = sorted(input, key=itemgetter(1))
>>> groups = groupby(sorted_input, key=itemgetter(1))
>>> [{'type':k, 'items':[x[0] for x in v]} for k, v in groups]
[{'items': ['5238761', '5349618', '962142', '7795297', '7341464', '5594916', '1550003'], 'type': 'ETH'}, {'items': ['11013331', '9843236'], 'type': 'KAT'}, {'items': ['9085267', '11788544'], 'type': 'NOT'}]

---

注意：在python 3.7之前，这两种方法都不保留键的原始顺序。如果你需要保持顺序，就得用OrderedDict。

>>> from collections import OrderedDict
>>> res = OrderedDict()
>>> for v, k in input:
...   if k in res: res[k].append(v)
...   else: res[k] = [v]
... 
>>> [{'type':k, 'items':v} for k,v in res.items()]
[{'items': ['11013331', '9843236'], 'type': 'KAT'}, {'items': ['9085267', '11788544'], 'type': 'NOT'}, {'items': ['5238761', '5349618', '962142', '7795297', '7341464', '5594916', '1550003'], 'type': 'ETH'}]

在python 3.7及之后的版本中，普通的字典会保持插入的顺序。