根据标签分组

问题描述

假设我有两个列表：letters = [a,b,c,d,e,f,g,h,i]和digits = [0,1,2,0]。我希望最终结果为output = [[a,i],[b,[d,h]]。

这里0，1和2只是不同的类。例如，a来自类0，而b和c来自类1。我只需要根据字母的类别将字母放在子列表中即可。

我想我可以在这里zip()和列表理解，但是我不确定该怎么做。如何使用numpy做到这一点？

解决方法

您可以使用zip()和临时词典：

letters = ['a','b','c','d','e','f','g','h','i']
digits = [0,1,2,0]

tmp = {}
for d,l in zip(digits,letters):
    tmp.setdefault(d,[]).append(l)

out = []
for k in sorted(tmp):
    out.append(tmp[k])

print(out)

打印：

[['a','f'],['b','i'],['d','h']]

---

或者：另一个版本（使用itertools.groupby）：

from itertools import groupby

out = []
for _,g in groupby(sorted(zip(digits,letters)),lambda k: k[0]):
    out.append([l for _,l in g])

print(out)

,

numpy解决方案也是可能的：

digits = np.array([0,0])
letters = np.array(['a','i'])
argidx = np.argsort(digits)
digits,letters = digits[argidx],letters[argidx]
markers = np.diff(digits,prepend=digits[0])
marker_idx,= np.nonzero(markers)
np.split(letters,marker_idx)

输出：

[array(['a',dtype='<U1'),array(['b','e'],array(['d','h'],dtype='<U1')]

,

自从您标记了numpy（我个人更喜欢pandas groupby，因为它很干净，并且专门用于此目的）：

d,l = list(zip(*(sorted(zip(digits,letters)))))
d = np.array(d)
idx = np.flatnonzero(np.r_[True,d[:-1] != d[1:],True])
output = [list(l[i:j]) for i,j in zip(idx[:-1],idx[1:])]
#[['a','h']]

,

答案的另一种变化形式（此处为groupby和itemgetter）：

from itertools import groupby
from operator import itemgetter

letters = ['a',0]

combined_letters_digits = sorted(zip(letters,digits),key=itemgetter(1))
letter_groups = groupby(combined_letters_digits,itemgetter(1))
out = [[item[0] for item in group_data] for (key,group_data) in letter_groups]

print(out)

,

更简洁，pandas没有循环

import pandas as pd

s = pd.Series(
    ['a',index=[0,1],name='letters'
)

[*s.groupby(level=0).agg(list)]


[['a','h']]