如何用NumPy API实现列三元组按元素逻辑或的求和计算？

问题背景 #

给定形状为(n, m)的布尔型NumPy数组A：

import numpy as np
A = np.array([[ True, False, False, False],
              [ True, False, False,  True],
              [False,  True, False, False],
              [False, False, False,  True],
              [False, False, False,  True]])

需要生成形状为(m, m, m)的整型数组R，其中R[k, i, j]的计算规则为：对A的第k、i、j列逐元素执行逻辑或操作后，统计结果为True的元素个数。示例计算如下：

R[3, 1, 2] = (np.logical_or(np.logical_or(A[:, 3], A[:, 1]), A[:, 2])).sum()

当前已有嵌套循环的实现方式，但未利用NumPy的向量化特性，现需通过NumPy广播机制实现该功能。

向量化实现方案 #

可以通过广播机制将列向量扩展为三维数组，一次性完成逻辑或与求和操作，完全避免循环，代码如下：

# 转置数组，将原数组的列转为行，形状变为(m, n)
A_T = A.T
# 利用广播生成三个维度的扩展数组，逐元素逻辑或后沿最后一维求和
R = (A_T[:, None, None] | A_T[None, :, None] | A_T[None, None, :]).sum(axis=-1)

方案细节解释 #

1. 数组转置：将A转置为A_T，形状从(n, m)变为(m, n)，此时A_T的每一行对应原数组A的一列，便于后续按维度扩展。
2. 广播扩展维度：
   • A_T[:, None, None]：将A_T扩展为(m, 1, 1, n)，对应结果数组R中的k维度
   • A_T[None, :, None]：扩展为(1, m, 1, n)，对应R中的i维度
   • A_T[None, None, :]：扩展为(1, 1, m, n)，对应R中的j维度
     这三个数组会通过NumPy的广播机制自动对齐为(m, m, m, n)的形状，实现所有k、i、j组合的列元素对齐。
3. 逻辑或与求和：对对齐后的数组逐元素执行逻辑或操作，然后沿最后一个维度（即原数组的行维度n）求和，最终得到形状为(m, m, m)的结果数组R，与需求完全匹配。

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