可以矢量化切片的最小值和最大值吗？

Question

假设我有一个 NumPy 整数数组。

arr = np.random.randint(0, 1000, 1000)

我有两个数组lower 和upper，分别代表arr 切片的下限和上限。这些间隔是重叠的并且是可变长度的，但是lowers 和uppers 都保证不会减少。

lowers = np.array([0, 5, 132, 358, 566, 822])
uppers = np.array([45, 93, 189, 533, 800, 923])

我想找到由lowers 和uppers 定义的arr 的每个切片的最小值和最大值，并将它们存储在另一个数组中。

out_arr = np.empty((lowers.size, 2))

最有效的方法是什么？我担心没有矢量化方法，因为我看不到如何在循环中绕过索引..

---

我目前的方法很简单

for i in range(lowers.size):
    arr_v = arr[lowers[i]:uppers[i]]
    out_arr[i,0] = np.amin(arr_v)
    out_arr[i,1] = np.amax(arr_v)

这让我得到了想要的结果，比如

In [304]: out_arr
Out[304]: 

array([[  26.,  908.],
       [  18.,  993.],
       [   0.,  968.],
       [   3.,  999.],
       [   1.,  998.],
       [   0.,  994.]])

但这对我的实际数据来说太慢了。

Answer 1

好的，下面是如何使用np.minimum.reduceat 至少缩小原始问题的大小：

lu = np.r_[lowers, uppers]
so = np.argsort(lu)
iso = np.empty_like(so)
iso[so] = np.arange(len(so))
cut = len(lowers)
lmin = np.minimum.reduceat(arr, lu[so])
for i in range(cut):
    print(min(lmin[iso[i]:iso[cut+i]]), min(arr[lowers[i]:uppers[i]]))

# 33 33
# 7 7
# 5 5
# 0 0
# 3 3
# 7 7

这并没有实现摆脱主循环，但至少数据从 1000 个元素的数组减少到 12 个元素的数组。

更新：

@Eric Hansen 自己的解决方案重叠很小。我仍然想指出，如果存在大量重叠，那么将这两种方法结合起来甚至可能是值得的。我没有numba，所以下面只是一个双通道版本，它结合了我的预处理和 Eric 的纯numpy 解决方案，它也可以作为onepass 形式的基准：

import numpy as np
from timeit import timeit

def twopass(lowers, uppers, arr):
    lu = np.r_[lowers, uppers]
    so = np.argsort(lu)
    iso = np.empty_like(so)
    iso[so] = np.arange(len(so))
    cut = len(lowers)
    lmin = np.minimum.reduceat(arr, lu[so])
    return np.minimum.reduceat(lmin, iso.reshape(2,-1).T.ravel())[::2]

def onepass(lowers, uppers, arr):
    mixture = np.empty((lowers.size*2,), dtype=lowers.dtype) 
    mixture[::2] = lowers; mixture[1::2] = uppers
    return np.minimum.reduceat(arr, mixture)[::2]

arr = np.random.randint(0, 1000, 1000)
lowers = np.array([0, 5, 132, 358, 566, 822])
uppers = np.array([45, 93, 189, 533, 800, 923])

print('small')
for f in twopass, onepass:
    print('{:18s} {:9.6f} ms'.format(f.__name__, 
                                     timeit(lambda: f(lowers, uppers, arr),
                                            number=10)*100))

arr = np.random.randint(0, 1000, 10**6)
lowers = np.random.randint(0, 8*10**5, 10**4)
uppers = np.random.randint(2*10**5, 10**6, 10**4)
swap = lowers > uppers
lowers[swap], uppers[swap] = uppers[swap], lowers[swap]


print('large')
for f in twopass, onepass:
    print('{:18s} {:10.4f} ms'.format(f.__name__, 
                                     timeit(lambda: f(lowers, uppers, arr),
                                            number=10)*100))

示例运行：

small
twopass             0.030880 ms
onepass             0.005723 ms
large
twopass               74.4962 ms
onepass             3153.1575 ms

Answer 2

执行会很慢，因为在循环内部，子数组被复制到数组中，然后执行操作。您可以通过单行代码避免整个循环

out_array = np.array([(np.amin(arr[lowers[i]:uppers[i]]),np.amax(arr[lowers[i]:uppers[i]])) for i in range(lowers.size)])

Answer 3

根据 Paul Panzer 对reduceat 的建议，我最初尝试的改进版本是

mixture = np.empty((lowers.size*2,), dtype=lowers.dtype) 
mixture[::2] = lowers; mixture[1::2] = uppers

np.column_stack((np.minimum.reduceat(arr, mixture)[::2],
                 np.maximum.reduceat(arr, mixture)[::2]))

在与我的实际数据相当的样本量上，这在我的机器上运行时间为 4.22 毫秒，而我的原始解决方案需要 73 毫秒。

在我原来的解决方案中使用 Numba 会更快

from numba import jit

@jit
def get_res():
    out_arr = np.empty((lowers.size, 2))
    for i in range(lowers.size):
        arr_v = arr[lowers[i]:uppers[i]]
        out_arr[i,0] = np.amin(arr_v)
        out_arr[i,1] = np.amax(arr_v)
    return out_arr

在我的机器上运行 100 微秒。