Python高效比对嵌套列表提取ID：百万级数据优化方案求助

问题背景 #

列表a包含约20万条四元组数据，列表b包含约100万条五元组数据（最后一位为ID）。需要提取列表b中前四个元素与列表a完全匹配的ID，原双重循环方法可行但处理百万级数据时速度极慢，需优化实现。

原代码示例 #

#      g    f    z   n
a = [(4264, 7, 1526, 0),
     (4293, 14, 846, 93),
     (4174, 6, 962, 0),
     (4256, 12, 121, 0),
     (4257, 29, 182, 385)
     ]

# list a ca. 200,000 entries

#      g    f    z   n  id
b = [(4264, 10, 397, 0, 113),
     (4264, 20, 95, 0, 114),
     (4279, 13, 41, 0, 115),
     (4293, 14, 846, 93, 116),
     (4264, 8, 94, 0, 117),
     (4264, 8, 92, 0, 118),
     (4256, 12, 121, 0,119),
     (4264, 9, 293, 82, 120),
     (4264, 9, 288, 0, 121),
     (4264, 8, 90, 25, 122),
     (4264, 9, 156, 0, 123)
     ]

# list b ca. 1,000,000 entries

# 原方法可行但速度极慢
xyz = []
for i in a:
        for x in b:
                if i[0] == x[0] and i[1] == x[1] and i[2] == x[2] and i[3] == x[3]:
                        xyz.append(x[4])

print(xyz)

# 预期结果: [116, 119]

优化方案 #

方案一：集合快速查找（推荐首选）

利用集合O(1)的查找特性，先将列表a的所有四元组存入集合，再遍历列表b判断前四个元素是否在集合中，时间复杂度降至O(len(a)+len(b))。

# 预处理a为集合，集合查找效率远高于线性遍历
a_set = set(a)

xyz = []
for item in b:
    # 取当前元素前四位组成元组，判断是否在a_set中
    if item[:4] in a_set:
        xyz.append(item[4])

print(xyz)  # 输出: [116, 119]

方案二：字典映射预处理

如果需要根据a的元素快速查找对应所有ID（比如b中存在多个ID对应同一前四元组的情况），可以先将列表b预处理为字典，键是前四元组，值是对应ID的列表。

# 构建b的前四元组到ID的映射字典
b_map = {}
for item in b:
    key = item[:4]
    # 若键不存在则初始化空列表，再追加ID
    b_map.setdefault(key, []).append(item[4])

xyz = []
for key in a:
    # 若a中的四元组存在于映射中，扩展结果列表
    if key in b_map:
        xyz.extend(b_map[key])

print(xyz)  # 输出: [116, 119]

方案三：Pandas批量处理（超大规模数据场景）

当数据量达到千万级时，使用Pandas的合并操作可以充分利用其内部优化的算法，处理效率更高。

import pandas as pd

# 将列表转换为DataFrame
df_a = pd.DataFrame(a, columns=['g', 'f', 'z', 'n'])
df_b = pd.DataFrame(b, columns=['g', 'f', 'z', 'n', 'id'])

# 按前四列合并两个DataFrame，只保留匹配的行
merged_df = pd.merge(df_a, df_b, on=['g', 'f', 'z', 'n'], how='inner')

# 提取ID列并转换为列表
xyz = merged_df['id'].tolist()

print(xyz)  # 输出: [116, 119]

方案选择建议 #

• 方案一：最简洁高效，内存占用低，适合大多数常规场景。
• 方案二：适合需要一对多映射的场景，方便快速获取同一匹配条件下的所有ID。
• 方案三：适合超大规模数据处理，熟悉Pandas的用户优先选择。

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