Python：如何高效剔除列表列表中各元素分属不同分组的组合

问题背景 #

生成1到52中选取6个数字的所有组合（列表的列表形式），需要剔除其中每个元素分别属于不同数字分组的组合。分组定义为D、T、L、H、K、M。原代码试图用if语句枚举6! = 720种排列情况来实现剔除，既繁琐又不现实，需要更简单高效的实现方式。

原代码 #

import numpy as np
import math
from itertools import permutations
import itertools
import collections


perm = itertools.combinations(range(1, 53), 6)
res = [list(val) for val in perm]
size = len(res)

D = [1, 2, 3, 4, 5, 6, 7, 8, 9]
T = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
L = [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
H = [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
K = [40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
M = [50, 51, 52]


output_list = []

for i in range(0, 20358520):
        if (
                res[i][0] not in D
                and res[i][1] not in T
                and res[i][2] not in L
                and res[i][3] not in H
                and res[i][4] not in K
                and res[i][5] in M
            )
            or 
            (
                res[i][0] not in D
                and res[i][2] not in T
                and res[i][1] not in L
                and res[i][3] not in H
                and res[i][4] not in K
                and res[i][5] not in M
            )
            or ... to 720 combinations:
                            
           output_list.append([res[i][0],
                               res[i][1],
                               res[i][2],
                               res[i][3],
                               res[i][4],
                               res[i][5]]
                            )

print(output_list )
print(len(output_list))

优化方案 #

核心思路：先建立数字到分组的映射字典实现快速查询，再通过判断组合的分组集合是否覆盖全部6个分组，直接筛选出需要保留的组合。

优化后代码

import itertools

# 定义分组
groups = {
    'D': [1,2,3,4,5,6,7,8,9],
    'T': [10,11,12,13,14,15,16,17,18,19],
    'L': [20,21,22,23,24,25,26,27,28,29],
    'H': [30,31,32,33,34,35,36,37,38,39],
    'K': [40,41,42,43,44,45,46,47,48,49],
    'M': [50,51,52]
}

# 构建数字到分组的映射字典，实现O(1)查询效率
num_to_group = {}
for group_name, nums in groups.items():
    for num in nums:
        num_to_group[num] = group_name

all_group_names = set(groups.keys())  # 所有分组的集合：{'D','T','L','H','K','M'}
output_list = []

# 直接遍历组合生成器，无需预存所有组合，大幅节省内存
for combo in itertools.combinations(range(1,53), 6):
    # 获取当前组合中所有数字对应的分组集合
    combo_groups = {num_to_group[num] for num in combo}
    # 如果分组集合不等于全部6个分组，说明不是每个元素分属不同分组，保留该组合
    if combo_groups != all_group_names:
        output_list.append(list(combo))

print(len(output_list))
# print(output_list)  # 如需查看具体结果可取消注释

方案优势

1. 查询高效：用字典存储数字到分组的映射，查询时间复杂度为O(1)，远快于原代码中num in list的O(n)查询。
2. 逻辑简洁：通过集合对比直接判断是否是需要剔除的组合，无需枚举720种排列情况。
3. 内存优化：直接遍历itertools.combinations生成器，不用一次性把所有组合存入列表，避免了大量内存占用。

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