统计字典键在列表列表中的出现次数并求argmax的高效实现

输入数据结构 #

原始代码示例：

import numpy as np
list_of_list = [
    ['ccc', 'cccc', 'b', 'c', 'b'],
    ['ab', 'b', 'b', 'aa'], 
    ['c', 'b', 'c', 'c', 'b', 'c'],
    ['bb', 'd', 'c'],
]
my_dict = {key: None for key in 'abcde'} 

其中list_of_list是由词汇子列表组成的列表，通用格式为：

list_of_list = [ 
    ['word1', 'word2', ... , 'wordN'],
    ['word1', 'word2', ... , 'wordM'], 
    ['word1', 'word2', ... , 'wordK'], 
    ...
]

需求目标 #

更新my_dict，让每个键对应的值格式为[最大出现次数的子列表索引, 最大出现次数]，即：

"key": [index_of_max_occurrence, max_occurrence]

若该键在所有子列表中均未出现，则保留None。

现有低效方案 #

以下嵌套循环的代码在小数据量下可用，但数据量大时性能极差：

for k in my_dict:
    counters = list()
    for lst in list_of_list:
        counters.append( lst.count(k) )
    if any(counters):
        my_dict[k] = [ np.argmax(counters) , max(counters) ]
print(my_dict) # {'a': None, 'b': [0, 2], 'c': [2, 4], 'd': [3, 1], 'e': None}

高效优化方案 #

核心思路 #

原方案的问题在于对每个键都重复遍历所有子列表并调用count，时间复杂度为O(MNK)（M为键的数量，N为子列表数量，K为子列表平均长度）。优化方向是一次性预统计所有子列表的元素频次，再基于统计结果生成目标字典，时间复杂度可降至O(N*K + M)。

基础优化实现 #

import numpy as np
from collections import defaultdict

list_of_list = [
    ['ccc', 'cccc', 'b', 'c', 'b'],
    ['ab', 'b', 'b', 'aa'], 
    ['c', 'b', 'c', 'c', 'b', 'c'],
    ['bb', 'd', 'c'],
]
my_dict = {key: None for key in 'abcde'} 

# 预统计：{元素: [(子列表索引, 出现次数), ...]}
element_counts = defaultdict(list)
for idx, sub_list in enumerate(list_of_list):
    # 统计当前子列表的元素频次
    sub_counter = defaultdict(int)
    for item in sub_list:
        sub_counter[item] += 1
    # 将当前子列表的统计结果存入全局字典
    for item, cnt in sub_counter.items():
        element_counts[item].append( (idx, cnt) )

# 填充目标字典
for key in my_dict:
    if key not in element_counts:
        continue
    # 找到当前元素的最大出现次数记录
    max_entry = max(element_counts[key], key=lambda x: x[1])
    my_dict[key] = [max_entry[0], max_entry[1]]

print(my_dict)  # {'a': None, 'b': [0, 2], 'c': [2, 4], 'd': [3, 1], 'e': None}

超大数据量进阶优化 #

如果子列表数量或元素规模极大，可结合pandas进行分组统计，利用其向量化运算进一步提升效率：

import pandas as pd

list_of_list = [
    ['ccc', 'cccc', 'b', 'c', 'b'],
    ['ab', 'b', 'b', 'aa'], 
    ['c', 'b', 'c', 'c', 'b', 'c'],
    ['bb', 'd', 'c'],
]
my_dict = {key: None for key in 'abcde'} 

# 构造包含元素和子列表索引的数据集
data = []
for idx, sub_list in enumerate(list_of_list):
    data.extend( [(item, idx) for item in sub_list] )
df = pd.DataFrame(data, columns=['item', 'list_idx'])

# 统计每个元素在各子列表中的出现次数
count_df = df.groupby(['item', 'list_idx']).size().reset_index(name='count')

# 筛选每个元素的最大次数记录
max_count_df = count_df.loc[count_df.groupby('item')['count'].idxmax()]

# 填充目标字典
for key in my_dict:
    row = max_count_df[max_count_df['item'] == key]
    if not row.empty:
        my_dict[key] = [row['list_idx'].values[0], row['count'].values[0]]

print(my_dict)

性能优势说明 #

1. 避免重复遍历：原方案每个键都要遍历所有子列表，优化后仅需遍历一次所有子列表完成预统计。
2. 子列表内统计更高效：使用defaultdict(int)统计子列表元素频次，只需遍历子列表一次；而原方案的lst.count(k)每次都要遍历子列表，存在大量重复计算。

内容的提问来源于stack exchange，提问作者farid