NumPy中筛选指定矩形范围内2D向量的更简洁高效实现及性能差异问询

我最近在处理一个需求：从NumPy的2D数组里筛选出所有完全落在指定矩形区域内的2D向量——要求向量的每个元素都大于左下角边界inf，同时小于右上角边界sup。

我一开始写了这样一段代码实现需求：

import numpy as np

a = np.arange(12).reshape(6, 2)
inf = np.array([2, 2])
sup = np.array([9, 9])

b = (inf < a) & (a < sup)
r = a[b[:, 0] & b[:, 1]]

这段代码能正确得到结果：数组a是[[ 0 1] [ 2 3] [ 4 5] [ 6 7] [ 8 9] [10 11]]，筛选后的子数组r为[[4 5] [6 7]]，完全符合预期。但总觉得这段代码有点啰嗦，应该有更简洁的写法，所以想问问大家有没有更好的实现建议？

有意思的是，后来我看到了一个更紧凑的实现方式（来自furas的方案），但实际测试下来，它的执行速度反而比我最初的“丑写法”要慢。

我的测试环境是AMD Ryzen 7 3800X CPU + Python 3.12.7，测试代码如下：

import numpy as np, timeit as ti

t = 'i2'
a = np.random.randint(1024, size=(1000, 2)).astype(t)
inf = np.array([100, 200]).astype(t)
sup = np.array([300, 400]).astype(t)

def f0(a, inf, sup):
    return a[((inf < a) & (a < sup)).all(axis=1)]

def f1(a, inf, sup):
    b = (inf < a) & (a < sup)
    return a[b[:, 0] & b[:, 1]]

print(f'Minimum, median and maximum execution time in us:')
for fun in ('f0(a, inf, sup)', 'f1(a, inf, sup)'):
    t = 10**6 * np.array(ti.repeat(stmt=fun, setup=fun, globals=globals(), number=1, repeat=999))
    print(f'{fun:20} {np.amin(t):8,.3f} {np.median(t):8,.3f} {np.amax(t):8,.3f}')

测试结果如下：

Minimum, median and maximum execution time in us:
f0(a, inf, sup)       29.485    29.806   133.301
f1(a, inf, sup)       17.773    17.964    22.663

可以看到，更简洁的f0函数在最小、中位数和最大执行时间上都比f1慢不少，甚至最大时间差了好几倍。

想问问大家：有没有既简洁又高效的实现方式？另外也很好奇，为什么更紧凑的写法性能反而更差呢？

内容来源于stack exchange