如何编写根据数值约束获取任意维度numpy数组索引的通用函数

通用实现方案 #

你可以通过广播运算+布尔数组按行聚合的方式实现任意维度的通用约束判断，不需要手动拼接每一列的条件，代码如下：

import numpy as np

def get_valid_indices(arr, bounds):
    # 将边界转为numpy数组，适配广播运算
    min_bounds = np.array(bounds['min'])
    max_bounds = np.array(bounds['max'])
    # 逐元素判断每个维度是否在范围内，得到和arr同shape的布尔数组
    in_range = (arr >= min_bounds) & (arr <= max_bounds)
    # 按行聚合：只有一行的所有维度都满足条件才为True
    valid_mask = in_range.all(axis=1)
    # 返回符合条件的行索引
    return np.where(valid_mask)[0]

效果验证 #

原2维场景测试

# 原始测试数组
array = np.array([[ 1, 43],
       [81, 87],
       [92, 46],
       [71, 92],
       [18,  5],
       [98, 25],
       [94, 84],
       [82, 36],
       [75, 83],
       [76, 47]])
bounds = {'min': [5,20], 'max': [80,90]}
print(get_valid_indices(array, bounds))
# 输出：[8 9] 和手动拼接条件的结果一致

反转维度场景测试

array_reversed = np.roll(array,1,axis=1)
bounds_reversed = {'min': [20,5], 'max': [90,80]}
print(get_valid_indices(array_reversed, bounds_reversed))
# 输出：[8 9] 结果稳定正确

3维场景测试

# 生成10行3列的测试数组
arr_3d = np.random.randint(0,100,(10,3))
# 对应3个维度的边界约束
bounds_3d = {'min': [10,20,30], 'max': [70,80,90]}
print(get_valid_indices(arr_3d, bounds_3d))
# 可直接输出符合条件的行索引，不需要修改逻辑

原方法不稳定的原因 #

np.logical_and 和 np.bitwise_and 原生仅支持两个输入数组的运算，如果你直接传入超过两个条件会报错，就算嵌套调用如果顺序不对也容易出问题；如果需要用这两个方法实现多条件聚合，需要配合reduce使用：

# 等价的reduce写法
conditions = [(array[:,i] >= bounds['min'][i]) & (array[:,i] <= bounds['max'][i]) for i in range(array.shape[1])]
valid_mask = np.logical_and.reduce(conditions)

这种写法和前面的all聚合效果一致，但写法更繁琐，不如直接用广播+按行all的方式简洁稳定。

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