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 数组），所以我想知道是否有比上述 Cython 版本更 Pythonesque 和/或更有效的解决方案？

谢谢

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 可以加快最后的索引，特别是在处理许多这样的间隔/分组时 -

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