使用numpy digitize时，如何忽略数组中的掩码值或NaN？

这个坑我之前也踩过——np.digitize对掩码数组的支持确实不太贴心，它不会自动跳过掩码或NaN值，反而会给这些位置统一返回-1。这里有两个亲测好用的替代方案，你可以根据自己的场景选：

方案一：提取有效元素处理后回填（保留原数组结构） #

如果需要保留原数组的形状，只是把无效位置（NaN/掩码）的结果设为NaN（或你想要的标记），可以先过滤出有效元素，处理完再把结果放回去：

import numpy as np

# 示例数组和分箱边界
A = np.array([1.2, 2.5, np.nan, 4.1, 5.7])
bins = np.array([3.0, 6.0])

# 1. 找出非NaN的有效索引
valid_indices = ~np.isnan(A)
# 2. 对有效元素执行digitize
digitized_valid = np.digitize(A[valid_indices], bins)
# 3. 创建结果数组，有效位置填充结果，无效位置留NaN
result = np.full_like(A, np.nan)
result[valid_indices] = digitized_valid

print(result)
# 输出：array([0., 0., nan, 1., 1.])

这个方法逻辑清晰，不管原数组有多少无效值，都能精准保留结构，只处理有效数据。

方案二：用np.where一步到位（简洁版） #

如果不需要中间变量，想用一行代码搞定，可以结合np.where来直接替换无效位置的结果：

result = np.where(np.isnan(A), np.nan, np.digitize(A, bins))

这个方案本质和第一个逻辑一致，只是把过滤、处理、回填合并成了一步，代码更紧凑，适合简单场景。

为什么原来的掩码数组方法没用？ #

np.ma.array的掩码只是屏蔽数组的显示和部分计算，但np.digitize并没有针对掩码数组做特殊适配——它只会读取数组的原始数据部分，然后对掩码对应的位置返回-1，所以达不到“忽略”的效果。上面的两个方案才是真正跳过无效值进行处理的正确姿势。

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