如何优化满足行列和约束的大规模二进制Numpy数组精简算法？

给定形状为(n,m)的二维二进制Numpy数组（元素仅为0或1），需在保证每行元素和≥min_x、每列元素和≥min_y的前提下，移除尽可能多的1，得到含1最少的合规数组。原代码在小矩阵上运行正常，但无法处理10x10以上的矩阵，完全不适配20000x500的大规模矩阵，需要优化。

• 暴力递归+全量状态缓存：用data.tobytes()作为缓存键，大规模矩阵的状态量呈指数级爆炸，内存和时间完全无法承受。
• 逐元素枚举尝试：遍历所有可移除行和列的组合，每一步时间复杂度为O(R*C)，递归深度叠加后开销剧增。
• 频繁矩阵复制：new_data = data.copy()对大矩阵来说是极高的内存冗余开销。
• 重复计算行/列和：每次递归都重新计算行、列总和，存在大量冗余计算。

核心思路 #

1. 优先满足列的最小要求：为每列锁定至少min_y个1，选择时优先占用当前已保留1数最少的行，避免行的配额被过度占用。
2. 补充行的剩余需求：对已保留1数不足min_x的行，从原矩阵的1位置中补充，优先选择当前已保留1数最少的列，确保总1数最少。
3. 移除所有多余1：仅保留满足约束的必要1，其余原矩阵中的1全部置0。

优化后的代码 #

import numpy as np

def min_dataset_optimized(data, min_x, min_y):
    data = data.astype(np.int8)
    n, m = data.shape
    
    # 先校验原矩阵是否可行
    row_sums = data.sum(axis=1)
    col_sums = data.sum(axis=0)
    if np.any(row_sums < min_x) or np.any(col_sums < min_y):
        return None
    
    # 初始化结果矩阵为全0，后续标记需要保留的1
    result = np.zeros_like(data)
    
    # 步骤1：满足每列至少min_y个1的要求
    for col in range(m):
        # 获取该列中原矩阵为1的所有行
        candidate_rows = np.where(data[:, col] == 1)[0]
        # 按当前行已保留的1数排序，优先选保留数少的行
        sorted_rows = candidate_rows[np.argsort(result[candidate_rows].sum(axis=1))]
        # 标记前min_y个行为保留
        selected_rows = sorted_rows[:min_y]
        result[selected_rows, col] = 1
    
    # 步骤2：满足每行至少min_x个1的要求
    row_remaining = np.maximum(min_x - result.sum(axis=1), 0)
    for row in range(n):
        need = row_remaining[row]
        if need <= 0:
            continue
        # 获取该行中原矩阵为1且未被标记的列
        candidate_cols = np.where((data[row] == 1) & (result[row] == 0))[0]
        # 按当前列已保留的1数排序，优先选保留数少的列
        sorted_cols = candidate_cols[np.argsort(result[:, candidate_cols].sum(axis=0))]
        # 标记前need个列为保留
        selected_cols = sorted_cols[:need]
        result[row, selected_cols] = 1
    
    # 验证约束（可选，用于调试）
    assert np.all(result.sum(axis=1) >= min_x)
    assert np.all(result.sum(axis=0) >= min_y)
    
    return result

if __name__ == "__main__":
    # 测试示例矩阵
    data = np.array([
        [0, 1, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [1, 1, 0, 1, 1],
        [0, 1, 1, 0, 1],
        [0, 1, 1, 0, 1]
    ])
    print("原矩阵：")
    print(data)
    ans = min_dataset_optimized(data, 1, 1)
    print("\n优化后的解：")
    print(ans)
    print("总1数：", ans.sum())
    
    # 测试20000x500大规模矩阵
    print("\n测试20000x500矩阵：")
    large_data = np.random.randint(0, 2, (20000, 500))
    # 确保原矩阵满足基本约束（避免无解）
    row_deficit = large_data.sum(axis=1) < 5
    large_data[row_deficit] = np.random.randint(0, 2, (row_deficit.sum(), 500))
    col_deficit = large_data.sum(axis=0) < 200
    large_data[:, col_deficit] = np.random.randint(0, 2, (20000, col_deficit.sum()))
    
    import time
    start = time.time()
    large_ans = min_dataset_optimized(large_data, 5, 200)
    end = time.time()
    print(f"处理时间：{end - start:.2f}秒")
    print(f"总1数：{large_ans.sum()}")
    print(f"每行和≥5：{np.all(large_ans.sum(axis=1)>=5)}")
    print(f"每列和≥200：{np.all(large_ans.sum(axis=0)>=200)}")

性能说明 #

• 时间复杂度：O(nm)，主要来自两次行/列遍历，排序操作的规模远小于nm，可忽略。
• 内存开销：仅需存储原矩阵和结果矩阵，20000x500的矩阵仅占约10MB内存，完全可控。
• 正确性：该贪心策略能构造出接近理论最小总1数（max(n*min_x, m*min_y)）的解，受限于原矩阵结构时仍为可行范围内的最优解。

若需要严格的全局最优解（适用于中等规模矩阵），可将问题转化为二分图最小权匹配：

• 行、列作为二分图两侧节点，原矩阵中1的位置对应行到列的边，权值为1。
• 寻找满足每行至少min_x条边、每列至少min_y条边的最小权边集。
  但该方法时间复杂度为O((n+m)^3)，不适用于20000x500的超大规模矩阵，因此贪心策略是大规模场景的最优选择。

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