将元组列表按相同字段拆分为子列表

Question

我有一个很大的元组列表，里面的格式是这样的。每个元组的第二个字段是类别字段。

    [(1, 'A', 'foo'),
    (2, 'A', 'bar'),
    (100, 'A', 'foo-bar'),

    ('xx', 'B', 'foobar'),
    ('yy', 'B', 'foo'),

    (1000, 'C', 'py'),
    (200, 'C', 'foo'),
    ..]

那么，有什么最有效的方法可以把它们按类别（比如 A、B、C 等等）分成小列表呢？

Answer 1

如果你想从一堆元组中提取出多个单个元素的列表：

foo = ((1,2), (3, 4), (5, 6), (7,8) , (9, 10))
[[z[i] for z in foo] for i in (0,1)]

如果你更喜欢得到多个单个元素的元组：

zip(*[(1,4),(2,5),(3,6)])

Answer 2

collections.defaultdict

itertools.groupby 这个工具需要输入的数据先按照某个关键字段排好序。如果没有排好序，你就得先进行排序，这样会花费 O(n log n) 的时间。为了确保能在 O(n) 的时间内完成，你可以使用一个包含列表的 defaultdict：

from collections import defaultdict

dd = defaultdict(list)
for item in data:
    dd[item[1]].append(item)

res = list(dd.values())

print(res)

[[(1, 'A', 'foo'), (2, 'A', 'bar'), (100, 'A', 'foo-bar')],
 [('xx', 'B', 'foobar'), ('yy', 'B', 'foo')],
 [(1000, 'C', 'py'), (200, 'C', 'foo')]]

Answer 3

使用 itertools.groupby：

import itertools
import operator

data=[(1, 'A', 'foo'),
    (2, 'A', 'bar'),
    (100, 'A', 'foo-bar'),

    ('xx', 'B', 'foobar'),
    ('yy', 'B', 'foo'),

    (1000, 'C', 'py'),
    (200, 'C', 'foo'),
    ]

for key,group in itertools.groupby(data,operator.itemgetter(1)):
    print(list(group))

会得到

[(1, 'A', 'foo'), (2, 'A', 'bar'), (100, 'A', 'foo-bar')]
[('xx', 'B', 'foobar'), ('yy', 'B', 'foo')]
[(1000, 'C', 'py'), (200, 'C', 'foo')]

或者，如果你想把每个组放在一个子列表里，可以用列表推导式来实现：

[list(group) for key,group in itertools.groupby(data,operator.itemgetter(1))]

---

传给 itertools.groupby 的第二个参数是一个函数，这个函数会作用于 data（第一个参数）里的每个项目。这个函数应该返回一个 key。然后，itertools.groupby 会把所有相邻的、具有相同 key 的项目放在一起。

operator.itemgetter(1) 用来获取序列中的第二个项目。

举个例子，如果

row=(1, 'A', 'foo')

那么

operator.itemgetter(1)(row)

就等于 'A'。

---

正如 @eryksun 在评论中提到的，如果元组的类别是随机顺序的，那么在使用 itertools.groupby 之前，你必须先对 data 进行排序。这是因为 itertools.groupby 只会把相邻的、具有相同 key 的项目分到同一组。

要按类别对元组进行排序，可以使用：

data2=sorted(data,key=operator.itemgetter(1))