如何用NumPythonic方法统计二维NumPy数组每行首个指定值的元素数（含该值）？

针对你提出的需求——统计二维NumPy数组每行中首个特定值出现的位置（包含该值，计数从1开始），若该行无此值则返回行长度，这里提供两种高效的NumPythonic实现方式：

方法1：结合argmax与any的向量操作 #

这是最简洁高效的方案，利用布尔数组的特性直接完成计算：

import numpy as np

val = 2
arr = np.array([
       [2, 2, 1, 1, 0],
       [0, 3, 1, 0, 0],
       [0, 1, 2, 0, 1]
])

# 生成布尔掩码：标记每个元素是否等于目标值
mask = arr == val
# 获取每行第一个匹配值的索引（0起始）
first_match_idx = mask.argmax(axis=1)
# 判断每行是否存在目标值
has_target = mask.any(axis=1)
# 构造结果：存在则索引+1（转为1起始计数），否则返回行长度
result = np.where(has_target, first_match_idx + 1, arr.shape[1])

print(result)  # 输出：array([1, 5, 3])

原理说明 #

• mask.argmax(axis=1)：在每行中返回第一个True的索引（若全为False则返回0）；
• mask.any(axis=1)：检查每行是否存在至少一个True（即目标值是否存在）；
• np.where：根据条件替换结果，完美处理"无目标值时返回行长度"的场景。

方法2：利用cumsum定位首个匹配值 #

如果需要更直观地追踪匹配值的出现顺序，可以使用累加操作：

mask = arr == val
# 对每行的布尔掩码做累加，首个1的位置就是目标值首次出现的位置
cumulative_mask = mask.cumsum(axis=1)
# 找到每行中第一个等于1的位置，转为1起始计数
first_match_pos = (cumulative_mask == 1).argmax(axis=1) + 1
# 对无目标值的行，替换为行长度
result = np.where(cumulative_mask.any(axis=1), first_match_pos, arr.shape[1])

print(result)  # 输出：array([1, 5, 3])

原理说明 #

• mask.cumsum(axis=1)：每行的布尔值累加后，首次出现目标值的位置会变为1，后续匹配值会累加为更大的数；
• (cumulative_mask == 1).argmax(axis=1)：精准定位到首个1的位置，再转为1起始计数。

两种方法均为纯向量操作，避免了Python循环，在处理大规模数组时能保持极高的效率。

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