Python中求取两列表交换N个元素所有可能组合的高效实现方法

核心思路 #

因为你的场景中N取值仅为1~4，完全不需要生成全量组合，仅需要分两步取组合再做替换即可：

• 从list1中选出N个需要被替换出去的元素组合
• 从list2中选出N个需要替换进来的元素组合
• 对每一组上述两个组合做元素替换，即可得到所有符合要求的结果，时间复杂度仅由N决定，和list1、list2的长度增长是低阶相关，性能远高于全组合方案。

Python原生实现代码 #

import itertools

def swap_n_elements(list1, list2, n):
    # 校验参数合法性
    if n < 0 or n > min(len(list1), len(list2)):
        return []
    res = []
    # 遍历所有从list1取n个元素的组合
    for remove_group in itertools.combinations(list1, n):
        remove_set = set(remove_group)
        # 保留list1中不需要替换的元素
        list1_remain = [x for x in list1 if x not in remove_set]
        # 遍历所有从list2取n个元素的组合
        for add_group in itertools.combinations(list2, n):
            add_set = set(add_group)
            # 保留list2中不需要替换的元素
            list2_remain = [x for x in list2 if x not in add_set]
            # 拼接得到新的两个列表
            new_list1 = list1_remain + list(add_group)
            new_list2 = list2_remain + list(remove_group)
            res.append( (new_list1, new_list2) )
    return res

测试验证 #

用你给出的示例测试：

list1 = [0,1]
list2 = [2,3]
print(swap_n_elements(list1, list2, 1))
# 输出 [([0, 2], [1, 3]), ([0, 3], [1, 2]), ([1, 2], [0, 3]), ([1, 3], [0, 2])] 与你的示例匹配
print(swap_n_elements(list1, list2, 2))
# 输出 [([2, 3], [0, 1])] 与你的示例匹配

量子化学场景适配 #

你提到的生成N阶激发Slater行列式的需求，正好对应上述逻辑：占据轨道集合对应list1，虚轨道集合对应list2，N就是激发阶数。如果你使用pyscf，可以直接调用pyscf.fci.addons.excitation_strings函数直接生成对应激发阶数的组态字符串，无需自己实现组合逻辑，适配性更高。

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