向量化查找分割 numpy 子数组最大值的过程

Question

给定一个 1d numpy 数组

arr

和一个索引数组

ind

,

numpy.split(arr, ind)

返回

arr

被

ind

划分的子数组列表。我们可以使用 Python

for

循环来获取所有子数组的最大值。

例如，

import numpy as np
arr = np.arange(12)
ind = np.array([3, 5, 9])
sub_arrays = np.split(arr, ind) 
# [array([0, 1, 2]), array([3, 4]), array([5, 6, 7, 8]), array([ 9, 10, 11])]
max_values = [sub_array.max() for sub_array in sub_arrays]
# [2, 4, 8, 11]

我们可以用 numpy 风格的矢量化方式来做到这一点吗？例如，一开始不拆分

arr

。

Answer 1

您可以使用

np.maximum.reduceat

。要获得准确的结果，您需要

ind

还包括索引 0。

import numpy as np

arr = np.arange(12)
ind = np.array([3, 5, 9])
max_values = np.maximum.reduceat(arr, np.concatenate(([0], ind)))
print(max_values) # [ 2  4  8 11]

Answer 2

import numpy as np

arr = np.arange(12)
ind = np.array([3, 5, 9])

# Create indices for reduction
indices = np.concatenate(([0], ind, [arr.size]))

# Use np.maximum.reduceat to find maximum values
max_values = np.maximum.reduceat(arr, indices[:-1])

print(max_values)