组argmax/argmin超过numpy中的分区索引 [英] group argmax/argmin over partitioning indices in numpy

问题描述

Numpy的ufunc具有 reduceat 方法，该方法可在数组中的连续分区上运行它们.所以不用写:

Numpy's ufuncs have a reduceat method which runs them over contiguous partitions within an array. So instead of writing:

import numpy as np
a = np.array([4, 0, 6, 8, 0, 9, 8, 5, 4, 9])
split_at = [4, 5]
maxima = [max(subarray for subarray in np.split(a, split_at)]

我可以写:

maxima = np.maximum.reduceat(a, np.hstack([0, split_at]))

两者都将返回切片a[0:4]，a[4:5]，a[5:10](即[8, 0, 9])中的最大值.

Both will return the maximum values in slices a[0:4], a[4:5], a[5:10], being [8, 0, 9].

我希望执行类似的功能 argmax ，请注意，我只希望每个分区中的一个单个最大索引:[3, 4, 5]具有上面的a和split_at(尽管索引5和9都在最后一组)，将由

I would like a similar function to perform argmax, noting that I would only like a single maximum index in each partition: [3, 4, 5] with the above a and split_at (despite indices 5 and 9 both obtaining the maximum value in the last group), as would be returned by

np.hstack([0, split_at]) + [np.argmax(subarray) for subarray in np.split(a, split_at)]

我将在下面发布一个可能的解决方案，但希望看到一个矢量化的解决方案，而不会在组上创建索引.

I will post a possible solution below, but would like to see one that is vectorized without creating an index over groups.

推荐答案

受此问题的启发，我在

Inspired by this question, ive added argmin/max functionality to the numpy_indexed package. Here is what the corresponding test looks like. Note that the keys may be in any order (and of any kind supported by npi):

def test_argmin():
    keys   = [2, 0, 0, 1, 1, 2, 2, 2, 2, 2]
    values = [4, 5, 6, 8, 0, 9, 8, 5, 4, 9]
    unique, amin = group_by(keys).argmin(values)
    npt.assert_equal(unique, [0, 1, 2])
    npt.assert_equal(amin,   [1, 4, 0])

这篇关于组argmax/argmin超过numpy中的分区索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！