NumPy 数组的间隔 argmax，间隔由另一个数组定义

Question

我有一个一维数组，其中包含已排序的非唯一数字。他们重复的次数是随机的。它与具有相同大小的权重数组相关联。对于给定的相同元素系列，相关联的权重系列也可能有也可能没有重复元素，并且在整个权重数组中，可能有也可能没有重复元素。例如：

pos     = np.array([3, 3, 7, 7, 9, 9, 9, 10, 10])
weights = np.array([2, 10, 20, 8, 5, 7, 15, 7, 2])

我需要提取 pos 的唯一元素数组，但唯一元素是权重最大的元素。

我提出的工作解决方案涉及循环：

pos     = np.array([3, 3, 7, 7, 9, 9, 9, 10, 10])
weights = np.array([2, 10, 20, 8, 5, 7, 15, 7, 2])
# Get the number of occurences of the elements in pos but throw away the unique array, it's not the one I want.
_, ucounts = np.unique(pos, return_counts=True)
# Initialize the output array.
unique_pos_idx = np.zeros([ucounts.size], dtype=np.uint32)

last = 0
for i in range(ucounts.size):
    maxpos = np.argmax( weights[last:last+ucounts[i]] )
    unique_pos_idx[i] = last + maxpos
    last += ucounts[i]

# Result is:
# unique_pos_idx = [1 2 6 7]

但我没有使用很多 Python 语言或 Numpy（除了使用 numpy 数组）所以我想知道是否有更多 Pythonesque and/or甚至比上面的 Cython 版本更有效的解决方案？

谢谢

Answer 1

这是一种矢量化方法 -

sidx = np.lexsort([weights,pos])
out = sidx[np.r_[np.flatnonzero(pos[1:] != pos[:-1]), -1]]

可能的性能改进 -

1] 使用 scaling -

获取排序索引 sidx 的更快方法

sidx = (pos*(weights.max()+1) + weights).argsort()

2] 使用 boolean-indexing 可以使最后的索引更快，特别是在处理许多这样的 intervals/groupings -

时

out = sidx[np.concatenate((pos[1:] != pos[:-1], [True]))]

运行时测试

所有方法：

def org_app(pos, weights):
    _, ucounts = np.unique(pos, return_counts=True)
    unique_pos_idx = np.zeros([ucounts.size], dtype=np.uint32)    
    last = 0
    for i in range(ucounts.size):
        maxpos = np.argmax( weights[last:last+ucounts[i]] )
        unique_pos_idx[i] = last + maxpos
        last += ucounts[i]
    return unique_pos_idx

def vec_app(pos, weights):
    sidx = np.lexsort([weights,pos])
    return sidx[np.r_[np.flatnonzero(pos[1:] != pos[:-1]), -1]]

def vec_app_v2(pos, weights):
    sidx = (pos*(weights.max()+1) + weights).argsort()
    return sidx[np.concatenate((pos[1:] != pos[:-1], [True]))]

时间和验证 -

对于设置，让我们使用示例并通过缩放将其平铺 10000 次，因为我们打算创建 1000 倍的间隔数。另外，让我们在 weights 中使用唯一数字，这样 argmax 索引就不会被相同的数字混淆：

In [155]: # Setup input
     ...: pos = np.array([3, 3, 7, 7, 9, 9, 9, 10, 10,])
     ...: pos = (pos + 10*np.arange(10000)[:,None]).ravel()
     ...: weights = np.random.choice(10*len(pos), size=len(pos), replace=0)
     ...: 
     ...: print np.allclose(org_app(pos, weights), vec_app(pos, weights))
     ...: print np.allclose(org_app(pos, weights), vec_app_v2(pos, weights))
     ...: 
True
True

In [156]: %timeit org_app(pos, weights)
     ...: %timeit vec_app(pos, weights)
     ...: %timeit vec_app_v2(pos, weights)
     ...: 
10 loops, best of 3: 56.4 ms per loop
100 loops, best of 3: 14.8 ms per loop
1000 loops, best of 3: 1.77 ms per loop

In [157]: 56.4/1.77 # Speedup with vectorized one over loopy
Out[157]: 31.864406779661017

NumPy 数组的间隔 argmax，间隔由另一个数组定义

Intervaled argmax for a NumPy array with intervals defined by another array

arrays

numpy

unique

extraction

weighted