如何用NumPy向量化方法替代循环统计符合条件的矩阵组合数？

问题描述 #

给定两个仅含0和1的矩阵：

• A：形状为n×m，n行代表设施，m列代表产品（A[n,m]=1表示设施n提供产品m）
• B：形状为n×o，o列代表客户地点（B[n,o]=1表示设施n服务客户o）

需要统计所有客户-产品组合（共o×m个）中，至少被一个设施覆盖的组合总数（即存在设施n，使得A[n,m]和B[n,o]同时为1）。

向量化解决方案 #

原循环实现效率较低，可通过矩阵乘法直接完成向量化计算，核心思路是利用矩阵乘法的特性，一次性计算所有组合的覆盖情况：

import numpy as np

# 原矩阵定义
A = np.array(
    [[1., 0., 1., 1., 1., 0., 1., 1., 1., 1.],
    [0., 0., 1., 1., 0., 1., 1., 0., 0., 1.],
    [0., 1., 1., 0., 0., 1., 0., 0., 1., 1.],
    [0., 0., 0., 0., 0., 1., 0., 1., 0., 0.],
    [0., 0., 1., 1., 0., 0., 0., 0., 0., 1.]]
)

B = np.array(
    [[0., 1., 0., 0., 1., 0., 1., 1., 0., 1., 0., 1., 1., 0., 1., 0., 0., 0., 1., 1.],
    [1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0., 0., 1., 1.],
    [0., 1., 1., 0., 1., 0., 0., 1., 0., 0., 0., 1., 1., 1., 1., 0., 0., 1., 0., 0.],
    [0., 1., 1., 1., 1., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0., 1., 0., 0., 0.],
    [1., 1., 0., 1., 0., 0., 0., 0., 0., 1., 1., 1., 1., 0., 1., 0., 1., 1., 0., 1.]]
)

def count_conditions_vectorized(A: np.ndarray, B: np.ndarray):
    assert np.all((A == 0) | (A == 1))
    assert np.all((B == 0) | (B == 1))
    
    # 计算A转置与B的矩阵乘积，得到m×o的矩阵
    # 每个元素coverage[m,o] = sum(A[n,m] * B[n,o] for n in 0..n-1)
    coverage = A.T @ B
    
    # 统计所有覆盖数大于0的组合数量
    return np.sum(coverage > 0)

# 验证结果
result = count_conditions_vectorized(A, B)
print(result)  # 输出145
assert result == 145

思路解释 #

原循环的核心逻辑是：对每个客户（B的列），计算每个产品是否被至少一个服务该客户的设施提供，然后累加数量。而矩阵乘法A.T @ B直接完成了所有组合的计算：

1. A.T是m×n的矩阵，代表产品-设施的对应关系
2. 矩阵乘法后得到m×o的矩阵coverage，其中coverage[m,o]表示同时提供产品m并服务客户o的设施数量
3. 只要coverage[m,o] > 0，就说明该产品-客户组合被至少一个设施覆盖，统计所有这样的元素总数即可得到结果

这种方法利用了NumPy底层优化的矩阵乘法（基于BLAS/LAPACK），相比Python循环，在大矩阵场景下性能提升极为显著。

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