NumPy中计算n维数组每行非-1相邻元素的列索引距离

输入样例 #

给定待处理的NumPy数组如下：

import numpy as np

arr = np.array([['a', -1, -1, -1],
                [ -1,'b','c', -1],
                ['e', -1,'d','f']], dtype=object)

计算规则 #

逐行遍历数组，仅当某行存在2个及以上取值不为-1的元素时做统计：

• 仅统计相邻非-1元素对：即两个非-1元素之间没有其他非-1元素，间隔存在其他非-1元素的配对不统计
• 距离计算方式：配对中右侧元素的列索引减去左侧元素的列索引

实现代码 #

result = []
for row in arr:
    # 提取当前行所有非-1元素的列索引，按从左到右顺序排列
    valid_cols = np.where(row != -1)[0]
    if len(valid_cols) >= 2:
        # 相邻索引做差，直接得到符合要求的距离
        row_distances = np.diff(valid_cols).tolist()
        result.append(row_distances)
    else:
        # 有效元素不足2个，无统计结果
        result.append([])

print(result)

结果校验 #

运行代码后输出结果为[[], [1], [2, 1]]，和示例规则完全匹配：

• 第0行仅存在1个非-1元素'a'，无符合要求的配对，返回空列表
• 第1行非-1元素列索引为[1, 2]，差值为1，对应元素对('b','c')的距离
• 第2行非-1元素列索引为[0, 2, 3]，相邻差值分别为2、1，对应元素对('e','d')、('d','f')的距离；('e','f')中间存在非-1元素'd'，不会被统计，符合要求

注：该实现兼容任意列数的二维NumPy数组，无需针对列数做额外适配

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