如何用Numpy寻找数组中行的组合，使各列求和值相等？

很遗憾，NumPy并没有内置工具能直接搞定这个需求——你要找的是矩阵行的子集，让它们的列求和结果完全一致，这本质上是多维子集和问题，属于NP难范畴，没有通用的高效内置解法。不过当矩阵规模不大时，我们可以结合NumPy的数组操作来实现一个可行的方案。

实现思路 #

核心思路是遍历所有可能的非空行子集，对每个子集计算列和，再检查所有列和是否相等：

• 用二进制掩码表示每个子集（比如掩码1011对应选中第0、1、3行）
• 借助NumPy的布尔索引快速筛选出选中的行
• 计算选中行的列和，判断所有列值是否一致

代码示例 #

import numpy as np

# 定义你的目标矩阵
matrix = np.array([[0,0,0,1], [1,0,1,0], [1,1,0,0], [0,1,0,0]])

n_rows = matrix.shape[0]

# 遍历所有非空子集（从1到2^n_rows - 1，排除空集）
for mask in range(1, 2 ** n_rows):
    # 把掩码转换成布尔数组，标记哪些行被选中
    selected_mask = np.array([(mask >> i) & 1 for i in range(n_rows)], dtype=bool)
    # 计算选中行的列和
    col_sums = matrix[selected_mask].sum(axis=0)
    # 检查所有列和是否相等
    if np.all(col_sums == col_sums[0]):
        # 获取选中行的索引（从0开始计数）
        row_indices = np.where(selected_mask)[0]
        print(f"找到符合条件的行索引：{row_indices}")
        print("列和结果：", col_sums)

运行这段代码后，你会得到输出：

找到符合条件的行索引：[0 1 3]
列和结果： [1 1 1 1]

注意事项 #

• 如果矩阵行数较多（比如超过20行），遍历所有子集的方法会变得异常缓慢，因为子集数量是2^n级别的指数增长。这种情况下你需要用到回溯、剪枝等优化算法，不过NumPy本身没有提供这类实现，得自己编写逻辑。
• 如果你需要找的是列和等于某个特定值（而不是只要所有列和相等），可以把判断条件改成np.all(col_sums == target_value)即可。

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