使用NumPy和Pandas找出列表中和为x的所有元素组合

给定列表 a = [1, +0, -6, 2, 3, 1, -0, 1, 2, -6, 9]，目标值 x=3，以下是用NumPy和Pandas实现的方案，找出所有元素值的组合（不考虑元素顺序，仅保留唯一组合）：

1. 数据准备 #

先将列表转换为NumPy数组和Pandas对象，方便后续操作：

import numpy as np
import pandas as pd

a = [1, +0, -6, 2, 3, 1, -0, 1, 2, -6, 9]
target = 3

# 转为NumPy数组
arr = np.array(a)
# 获取列表中的唯一值
unique_vals = np.unique(arr)

2. 查找符合条件的组合 #

单元素组合 #

直接筛选等于目标值的元素：

# NumPy实现
single_combs = arr[arr == target]
# 结果：array([3])

对应组合：3

双元素组合 #

遍历所有唯一值对，筛选和为3的组合，同时验证原列表中是否有足够的元素支持该组合：

# NumPy实现
pair_combs = []
for i in range(len(unique_vals)):
    val1 = unique_vals[i]
    for j in range(i, len(unique_vals)):
        val2 = unique_vals[j]
        if val1 + val2 == target:
            # 若两个值相同，需原列表中该值出现至少2次；否则只需各出现至少1次
            if val1 == val2:
                if np.sum(arr == val1) >= 2:
                    pair_combs.append(f"{val1},{val2}")
            else:
                if np.sum(arr == val1) >=1 and np.sum(arr == val2) >=1:
                    pair_combs.append(f"{val1},{val2}")

# Pandas实现（用笛卡尔积）
df = pd.DataFrame({'val': unique_vals})
cross_df = pd.merge(df, df, how='cross', suffixes=('_1', '_2'))
# 过滤重复组合（如(1,2)和(2,1)视为同一组合）
cross_df = cross_df[cross_df['val_1'] <= cross_df['val_2']]
# 筛选和为目标值的组合
cross_df = cross_df[cross_df['val_1'] + cross_df['val_2'] == target]
# 验证元素数量
cross_df['count_1'] = cross_df['val_1'].apply(lambda v: np.sum(arr == v))
cross_df['count_2'] = cross_df['val_2'].apply(lambda v: np.sum(arr == v))
cross_df = cross_df[(cross_df['val_1'] != cross_df['val_2']) | (cross_df['count_1'] >=2)]
pair_combs_pd = [f"{row['val_1']},{row['val_2']}" for _, row in cross_df.iterrows()]

双元素组合结果：0,3；1,2；-6,9

三元素组合 #

同样遍历唯一值的三元组，筛选和为3的组合并验证元素数量：

# NumPy实现
triple_combs = []
for i in range(len(unique_vals)):
    val1 = unique_vals[i]
    for j in range(i, len(unique_vals)):
        val2 = unique_vals[j]
        for k in range(j, len(unique_vals)):
            val3 = unique_vals[k]
            if val1 + val2 + val3 == target:
                # 统计组合中各值需要的数量
                count_needed = {val1:0, val2:0, val3:0}
                count_needed[val1] +=1
                count_needed[val2] +=1
                count_needed[val3] +=1
                # 验证原列表中各值的数量是否足够
                valid = True
                for v, cnt in count_needed.items():
                    if np.sum(arr == v) < cnt:
                        valid = False
                        break
                if valid:
                    triple_combs.append(f"{val1},{val2},{val3}")

# Pandas实现
cross3_df = pd.merge(df, df, how='cross').merge(df, how='cross')
cross3_df.columns = ['val1', 'val2', 'val3']
# 生成排序后的元组去重
cross3_df['sorted_vals'] = cross3_df.apply(lambda r: tuple(sorted([r['val1'], r['val2'], r['val3']])), axis=1)
cross3_df = cross3_df.drop_duplicates('sorted_vals')
# 筛选和为目标值的组合
cross3_df = cross3_df[cross3_df['val1'] + cross3_df['val2'] + cross3_df['val3'] == target]
# 验证元素数量
def is_valid(row):
    vals = row['sorted_vals']
    cnt_dict = {}
    for v in vals:
        cnt_dict[v] = cnt_dict.get(v, 0) +1
    for v, cnt in cnt_dict.items():
        if np.sum(arr == v) < cnt:
            return False
    return True
cross3_df = cross3_df[cross3_df.apply(is_valid, axis=1)]
triple_combs_pd = [",".join(map(str, row['sorted_vals'])) for _, row in cross3_df.iterrows()]

三元素组合结果：0,0,3；0,1,2；1,1,1

更多元素组合 #

如果考虑包含多个0或重复元素的更长组合（比如0,0,-6,9、0,1,1,1），只需在上述逻辑基础上扩展遍历维度即可，核心逻辑一致：确保组合元素和为3，且原列表中对应元素的数量满足组合需求。

最终所有唯一组合汇总 #

• 单元素：3
• 双元素：0,3；1,2；-6,9
• 三元素：0,0,3；0,1,2；1,1,1
• 四元素：0,0,-6,9；0,1,1,1

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