查找作为列表存在的列元素的数据框索引的最快方法

Posted 2023-03-11

【中文标题】查找作为列表存在的列元素的数据框索引的最快方法

【英文标题】：Fastest way to find dataframe indexes of column elements that exist as lists

【发布时间】：2019-09-02 12:59:18

【问题描述】：

我有一个熊猫数据框，其中列值作为列表存在。每个列表都有几个元素，一个元素可以存在于几行中。一个示例数据框是：

X = pd.DataFrame([(1,['a','b','c']),(2,['a','b']),(3,['c','d'])],columns=['A','B'])

X = 
 A          B
0  1  [a, b, c]
1  2  [a, b]
2  3     [c, d]

我想找到与列表中的元素相对应的所有行，即数据框索引，并从中创建一个字典。在这里忽略 A 列，因为 B 列是感兴趣的列！所以元素 'a' 出现在索引 0,1 中，它给出了 'a':[0,1]。此示例数据框的解决方案是：

Y = 'a':[0,1],'b':[0,1],'c':[0,2],'d':[2]

我编写了一个运行良好的代码，并且可以得到结果。我的问题更多与计算速度有关。我的实际数据框有大约 350,000 行，“B”列中的列表最多可以包含 1,000 个元素。但目前代码运行了几个小时！我想知道我的解决方案是否效率很低。 任何以更快更有效的方式提供的帮助将不胜感激！ 这是我的解决方案代码：

import itertools
import pandas as pd
X = pd.DataFrame([(1,['a','b','c']),(2,['a','b']),(3,['c','d'])],columns=['A','B'])
B_dict = []
for idx,val in X.iterrows():
    B = val['B']
    B_dict.append(dict(zip(B,[[idx]]*len(B))))
    B_dict = [k: list(itertools.chain.from_iterable(list(filter(None.__ne__, [d.get(k) for d in B_dict])))) for k in set().union(*B_dict)]

print ('Result:',B_dict[0])

输出

Result: 'd': [2], 'c': [0, 2], 'b': [0, 1], 'a': [0, 1]

for 循环中最后一行的代码是从这里借用的：Combine values of same keys in a list of dicts 和 remove None value from a list without removing the 0 value

【问题讨论】：

您要求进行代码优化，这是代码审查的一部分，可以在 here

找到

感谢您的快速回复！我也在 Code Review 论坛上发帖。 codereview.stackexchange.com/questions/217288/…

【参考方案1】：

使用此方法展开您的列表：https://***.com/a/46740682/9177877

然后分组并应用列表：

idx = np.arange(len(X)).repeat(X['B'].str.len(), 0)
s = X.iloc[idx, ].assign(B=np.concatenate(X['B'].values))['B']
d = s.to_frame().reset_index().groupby('B')['index'].apply(list).to_dict()

# 'a': [0, 1], 'b': [0, 1], 'c': [0, 2], 'd': [2]

150,000 行的速度非常快：

# sample data
X = pd.DataFrame([(1,['a','b','c']),(2,['a','b']),(3,['c','d'])],columns=['A','B'])
df = pd.concat([X]*50000).reset_index(drop=True)

%%timeit
idx = np.arange(len(df)).repeat(df['B'].str.len(), 0)
s = df.iloc[idx, ].assign(B=np.concatenate(df['B'].values))['B']
d = s.to_frame().reset_index().groupby('B')['index'].apply(list).to_dict()

# 530 ms ± 46.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

【讨论】：

感谢克里斯的解决方案！你的方法比我原来的解决方案快得多。但是还有一种更快的方法，如 ALollz 的 defaultdict 方法所示。所以我接受了它作为解决方案。谢谢！

【参考方案2】：

我认为defaultdict 将在大约 1 分钟内工作：

from collections import defaultdict
from itertools import chain

dd = defaultdict(list)
for k,v in zip(chain.from_iterable(df.B.ravel()), df.index.repeat(df.B.str.len()).tolist()):
    dd[k].append(v)

输出：

defaultdict(list, 'a': [0, 1], 'b': [0, 1], 'c': [0, 2], 'd': [2])

---

X = pd.DataFrame([(1, ['a', 'b', 'c']*300), (2, ['a', 'b']*50),
                  (3, ['c', 'd']*34)], columns=['A', 'B'])
df = pd.concat([X]*150000).reset_index(drop=True)

%%timeit
dd = defaultdict(list)
for k,v in zip(chain.from_iterable(df.B.ravel()), df.index.repeat(df.B.str.len()).tolist()):
    dd[k].append(v)
#38.1 s ± 238 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
idx = np.arange(len(df)).repeat(df['B'].str.len(), 0)
s = df.iloc[idx, ].assign(B=np.concatenate(df['B'].values))['B']
d = s.to_frame().reset_index().groupby('B')['index'].apply(list).to_dict()
#1min 24s ± 458 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

【讨论】：

非常感谢您非常有见地的回答！ defaultdict 方法比这里建议的其他方法快得多。所以我把它标记为正确答案。谢谢！

@network_coder。循环两次的代码审查的答案比这还要快，可能是因为 .str.len() 调用很慢。虽然相同的一般原则。

【参考方案3】：

X = pd.DataFrame([(1,['a','b','c']),(2,['a','b']),(3,['c','d'])],columns=['A','B'])

df = X['B'].apply(pd.Series).T.unstack().reset_index().drop(columns = ['level_1']).dropna()
df.groupby(0)['level_0'].apply(list).to_dict()

我将您的 B 列设为自己的 DF，将其转置以使 Index 成为列，取消堆叠，然后完成清理。它看起来像：

df

 level_0 0
0   0    a
1   0    b
2   0    c
3   1    a
4   1    b
6   2    c
7   2    d

然后我按第 0 列分组，使其成为一个列表，然后是一个字典。

【讨论】：

感谢本的回复。这也很有用，但是当数据大小增加时，解决方案会变慢。我在这里发布的问题很好：codereview.stackexchange.com/questions/217288/…。在这种情况下，您的解决方案的计算效率与我接受的最快解决方案的 defaultdict 解决方案进行了比较。谢谢！