NumPy遍历n维数组各列非零元素 输出双连接双向配对

问题背景 #

假设存在如下NumPy ndarray：

array([[['a',0,0,0],
        [0,'b','c',0],
        ['e','d',0,0]]])

需要按以下规则提取元素的双连接（double connections）：

• 数组每一列对应一个时间瞬时，列内元素为该时间点发生的事件，0代表无事件。比如第一个时间瞬时a、e同时发生，第二个时间瞬时b、d同时发生，第三个时间瞬时仅c发生。
• 同一时间瞬时如果恰好存在2个非零元素，就判定二者存在双连接，需要输出正反两个顺序的连接对；如果列内非零元素数量不是2，直接跳过该列。

预期输出格式如下：

('a','e')
('e','a')
('b','d')
('d','b')

原始测试数组可通过以下代码复现（注：开头展示的三层嵌套数组为笔误，实际复现的是3行4列二维数组）：

import numpy as np
array = np.array([['a',0,0,0],
    [0,'b','c',0],
    ['e','d',0,0]],dtype=object)

实现思路 #

直接按列遍历处理即可，步骤如下：

1. 利用NumPy数组的.T转置属性，转置后数组的每一行对应原数组的每一列，直接遍历转置数组就能逐列处理
2. 对每一列的元素做过滤，剔除值为0的项，得到当前时间点的有效事件列表
3. 判断有效事件列表长度，只有长度恰好为2时，生成正反两个顺序的元组输出；其余长度的列直接跳过

可运行代码 #

import numpy as np

# 初始化测试数组
arr = np.array([['a',0,0,0],
    [0,'b','c',0],
    ['e','d',0,0]],dtype=object)

# 逐列处理
for col in arr.T:
    # 过滤非零有效事件
    valid_events = [elem for elem in col if elem != 0]
    # 仅处理恰好2个有效事件的列
    if len(valid_events) == 2:
        print((valid_events[0], valid_events[1]))
        print((valid_events[1], valid_events[0]))

运行代码后输出和预期完全一致：

('a', 'e')
('e', 'a')
('b', 'd')
('d', 'b')

如果后续需要支持单时间点超过2个事件的全连接配对，只需要把长度判断部分替换为对应排列组合逻辑即可，当前代码完全匹配题目给出的规则。

内容的提问来源于stack exchange，提问作者Song Calderone Zhang