如何实现更快的Python分组笛卡尔积对遍历算法？

我有两个数据集（lhs和rhs），各包含一个分组变量字段（还有其他变量）。需求是：

• 仅保留lhs与rhs分组变量值相同的笛卡尔积对，忽略仅单侧存在的分组对应的项
• 遍历配对时要执行自定义逻辑（示例中是单词长度比较，实际是一系列函数调用）
• 分组变量值由外部参数动态生成，且数据量可达百万级，速度是核心需求，同时不希望对数据排序

目前我用字典手动建索引的方案能得到正确结果，但想找更快的实现；试过numpy版本，结果正确但速度更慢。请问有更高效的策略吗？

现有正确实现（字典索引版） #

lhs_data = [{'lhs_id':'001', 'word':'clean'},
            {'lhs_id':'002', 'word':'stop'},
            {'lhs_id':'003', 'word':'those'}]
rhs_data = [{'rhs_id':'001', 'word':'quack'},
            {'rhs_id':'002', 'word':'step'},
            {'rhs_id':'003', 'word':'stir'},
            {'rhs_id':'004', 'word':'state'},
            {'rhs_id':'005', 'word':'storm'},
            {'rhs_id':'006', 'word':'thin'},
            {'rhs_id':'008', 'word':'thaw'},
            {'rhs_id':'009', 'word':'thorn'},
            {'rhs_id':'010', 'word':'thumb'},
            {'rhs_id':'011', 'word':'true'},
            {'rhs_id':'012', 'word':'trade'}]

# 创建分组变量与索引映射
lhs_allgroups = {}
rhs_allgroups = {}
for recidx in range(len(lhs_data)):
   lhs_group = lhs_data[recidx]['word'][0:2]
   if lhs_group in lhs_allgroups:
      lhs_allgroups[lhs_group].append(recidx)
   else:
      lhs_allgroups[lhs_group] = [recidx]
for recidx in range(len(rhs_data)):
   rhs_group = rhs_data[recidx]['word'][0:2]
   if rhs_group in lhs_allgroups:
      if rhs_group in rhs_allgroups:
         rhs_allgroups[rhs_group].append(recidx)
      else:
         rhs_allgroups[rhs_group] = [recidx]

# 遍历配对并执行逻辑
for lhs_group in lhs_allgroups:
   if lhs_group in rhs_allgroups:
      for lhs_idx in lhs_allgroups[lhs_group]:
         for rhs_idx in rhs_allgroups[lhs_group]:
            if len(lhs_data[lhs_idx]['word']) == len(rhs_data[rhs_idx]['word']):
               print(lhs_data[lhs_idx]['lhs_id'], lhs_data[lhs_idx]['word'],
                     rhs_data[rhs_idx]['rhs_id'], rhs_data[rhs_idx]['word'])

正确输出：

002 stop 002 step
002 stop 003 stir
003 those 009 thorn
003 those 010 thumb

较慢的Numpy实现 #

import numpy as np

lhs_data = np.array([('001', 'clean'),
                     ('002', 'stop'),
                     ('003', 'those')],
                    dtype=[('lhs_id', 'S3'), ('word', 'S10')])
rhs_data = np.array([('001', 'quack'),
                     ('002', 'step'),
                     ('003', 'stir'),
                     ('004', 'state'),
                     ('005', 'storm'),
                     ('006', 'thin'),
                     ('008', 'thaw'),
                     ('009', 'thorn'),
                     ('010', 'thumb'),
                     ('011', 'true'),
                     ('012', 'trade')],
                    dtype=[('rhs_id', 'S3'), ('word', 'S10')])

# 生成分组变量
lhs_grouping_var = np.array([row['word'][:2] for row in lhs_data])
rhs_grouping_var = np.array([row['word'][:2] for row in rhs_data])

# 获取公共分组
lhs_unique = np.unique(lhs_grouping_var)
rhs_unique = np.unique(rhs_grouping_var)
common = np.intersect1d(lhs_unique, rhs_unique)

# 遍历配对
for group in common:
   lhs_idx = np.where(lhs_grouping_var == group)
   rhs_idx = np.where(rhs_grouping_var == group)
   for l in lhs_idx[0]:
      for r in rhs_idx[0]:
         if len(lhs_data[l]['word']) == len(rhs_data[r]['word']):
            print(lhs_data[l]['lhs_id'], lhs_data[l]['word'], rhs_data[r]['rhs_id'], rhs_data[r]['word'])

高效优化策略 #

1. 优化字典索引构建（轻量无依赖） #

用collections.defaultdict简化代码，同时预取分组内的所有记录，减少循环内的字典查找开销：

from collections import defaultdict

lhs_data = [{'lhs_id':'001', 'word':'clean'},
            {'lhs_id':'002', 'word':'stop'},
            {'lhs_id':'003', 'word':'those'}]
rhs_data = [{'rhs_id':'001', 'word':'quack'},
            {'rhs_id':'002', 'word':'step'},
            {'rhs_id':'003', 'word':'stir'},
            {'rhs_id':'004', 'word':'state'},
            {'rhs_id':'005', 'word':'storm'},
            {'rhs_id':'006', 'word':'thin'},
            {'rhs_id':'008', 'word':'thaw'},
            {'rhs_id':'009', 'word':'thorn'},
            {'rhs_id':'010', 'word':'thumb'},
            {'rhs_id':'011', 'word':'true'},
            {'rhs_id':'012', 'word':'trade'}]

# 构建lhs分组索引
lhs_groups = defaultdict(list)
for idx, rec in enumerate(lhs_data):
    group = rec['word'][:2]
    lhs_groups[group].append(idx)

# 仅保留lhs存在的分组，构建rhs索引
rhs_groups = defaultdict(list)
for idx, rec in enumerate(rhs_data):
    group = rec['word'][:2]
    if group in lhs_groups:
        rhs_groups[group].append(idx)

# 遍历配对，预取分组内记录减少查找
for group in lhs_groups:
    if group not in rhs_groups:
        continue
    # 预取当前分组的所有记录
    lhs_recs = [lhs_data[i] for i in lhs_groups[group]]
    rhs_recs = [rhs_data[i] for i in rhs_groups[group]]
    for lhs_rec in lhs_recs:
        lhs_len = len(lhs_rec['word'])
        for rhs_rec in rhs_recs:
            if len(rhs_rec['word']) == lhs_len:
                print(lhs_rec['lhs_id'], lhs_rec['word'], rhs_rec['rhs_id'], rhs_rec['word'])

2. Pandas向量化分组（适合百万级结构化数据） #

Pandas底层用C实现分组与合并操作，比纯Python循环快一个数量级，适合结构化数据场景：

import pandas as pd

lhs_df = pd.DataFrame(lhs_data)
rhs_df = pd.DataFrame(rhs_data)

# 生成分组变量
lhs_df['group'] = lhs_df['word'].str[:2]
rhs_df['group'] = rhs_df['word'].str[:2]

# 过滤仅单侧存在的分组
common_groups = lhs_df['group'].unique()
rhs_df = rhs_df[rhs_df['group'].isin(common_groups)]

# 内连接自动生成同分组的笛卡尔积
merged = lhs_df.merge(rhs_df, on='group', how='inner')

# 执行自定义过滤逻辑
filtered = merged[merged['word_x'].str.len() == merged['word_y'].str.len()]

# 输出结果
for _, row in filtered.iterrows():
    print(row['lhs_id'], row['word_x'], row['rhs_id'], row['word_y'])

3. Numba JIT加速（保留Python逻辑，接近C速度） #

如果需要保留自定义循环逻辑，用Numba编译循环，大幅提升执行速度：

from numba import jit
import numpy as np
from collections import defaultdict

# 转成NumPy结构化数组
lhs_dtype = [('lhs_id', 'U3'), ('word', 'U10')]
rhs_dtype = [('rhs_id', 'U3'), ('word', 'U10')]

lhs_np = np.array([(d['lhs_id'], d['word']) for d in lhs_data], dtype=lhs_dtype)
rhs_np = np.array([(d['rhs_id'], d['word']) for d in rhs_data], dtype=rhs_dtype)

# 构建分组索引
lhs_groups = defaultdict(list)
for idx in range(len(lhs_np)):
    group = lhs_np[idx]['word'][:2]
    lhs_groups[group].append(idx)

rhs_groups = defaultdict(list)
for idx in range(len(rhs_np)):
    group = rhs_np[idx]['word'][:2]
    if group in lhs_groups:
        rhs_groups[group].append(idx)

# Numba编译配对逻辑
@jit(nopython=True)
def process_pairs(lhs_data, rhs_data, lhs_groups, rhs_groups):
    results = []
    for group in lhs_groups:
        if group not in rhs_groups:
            continue
        lhs_idxs = lhs_groups[group]
        rhs_idxs = rhs_groups[group]
        for l_idx in lhs_idxs:
            lhs_word = lhs_data[l_idx]['word']
            lhs_len = len(lhs_word)
            lhs_id = lhs_data[l_idx]['lhs_id']
            for r_idx in rhs_idxs:
                rhs_word = rhs_data[r_idx]['word']
                if len(rhs_word) == lhs_len:
                    rhs_id = rhs_data[r_idx]['rhs_id']
                    results.append((lhs_id, lhs_word, rhs_id, rhs_word))
    return results

# 执行并输出
results = process_pairs(lhs_np, rhs_np, lhs_groups, rhs_groups)
for res in results:
    print(res[0], res[1], res[2], res[3])

方案对比 #

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