由标签数组定义的组的 NumPy 最大值

Question

我有两个数组，一个是值列表，一个是与每个值对应的 ID 列表。有些 ID 有多个值。我想创建一个新数组，其中包含为每个 id 记录的最大值，该数组的长度等于唯一 id 的数量。

使用for循环的示例：

import numpy as np

values = np.array([5, 3, 2, 6, 3, 4, 8, 2, 4, 8])
ids = np.array([0, 1, 3, 3, 3, 3, 5, 6, 6, 6])

uniq_ids = np.unique(ids)
maximums = np.ones_like(uniq_ids) * np.nan

for i, id in enumerate(uniq_ids):
    maximums[i] = np.max(values[np.where(ids == id)])

print(uniq_ids)
print(maximums)

[0 1 3 5 6]
[5. 3. 6. 8. 8.]

是否可以对其进行矢量化以使其运行速度更快？我正在想象一种可以仅使用 NumPy 函数创建“最大值”数组的单行代码，但我无法想出任何可行的方法。

Answer 1

用于np.lexsort同时对两个列表进行排序：

\n

idx = np.lexsort([values, ids])\n

\n

最后一个唯一 ID 的索引由下式给出

\n

last = np.r_[np.flatnonzero(np.diff(ids[idx])), len(ids) - 1]\n

\n

您可以使用它来获取每组的最大值：

\n

values[idx[last]]\n

\n

这与 相同values[idx][last]，但速度更快，因为您只需要len(last)通过这种方式提取元素，而不是重新排列整个数组然后提取。

\n

请记住，这np.unique基本上是“做”sort以及flatnonzero“何时做” return_index=True。

\n

这是迄今为止所有解决方案的集合以及基准。我建议对下面实施的@bb1 的 pandas 配方进行修改。

\n

def Shannon(values, ids):\n    uniq_ids = np.unique(ids)\n    maxima = np.ones_like(uniq_ids) * np.nan\n    for i, id in enumerate(uniq_ids):\n        maxima[i] = np.max(values[np.where(ids == id)])\n    return maxima\n\ndef richardec(values, ids):\n    return [a.max() for a in np.split(values, np.arange(1, ids.shape[0])[(np.diff(ids) != 0)])]\n\ndef MadPhysicist(values, ids):\n    idx = np.lexsort([values, ids])\n    return values[idx[np.r_[np.flatnonzero(np.diff(ids[idx])), len(ids) - 1]]]\n\ndef PeptideWitch(values, ids):\n    return np.vectorize(lambda id: np.max(values[np.where(ids == id)]))(np.unique(ids))\n\ndef mathfux(values, ids):\n    idx = np.argsort(ids)\n    return np.maximum.reduceat(values[idx], np.r_[0, np.flatnonzero(np.diff(ids[idx])) + 1])\n\ndef bb1(values, ids):\n    return DataFrame({\'ids\': ids, \'vals\': values}).groupby(\'ids\')[\'values\'].max().to_numpy()\n\ndef bb1_modified(values, ids):\n    return Series(values).groupby(ids).max().to_numpy()\n

\n

values = np.array([5, 3, 2, 6, 3, 4, 8, 2, 4, 8])\nids = np.array([0, 1, 3, 3, 3, 3, 5, 6, 6, 6])\n\n%timeit Shannon(values, ids)\n42.1 \xc2\xb5s \xc2\xb1 561 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n%timeit richardec(values, ids)\n27.7 \xc2\xb5s \xc2\xb1 181 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n%timeit MadPhysicist(values, ids)\n19.3 \xc2\xb5s \xc2\xb1 268 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n%timeit PeptideWitch(values, ids)\n55.9 \xc2\xb5s \xc2\xb1 1.29 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n%timeit mathfux(values, ids)\n20.9 \xc2\xb5s \xc2\xb1 308 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n%timeit bb1(values, ids)\n865 \xc2\xb5s \xc2\xb1 34.6 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit bb1_modified(values, ids)\n537 \xc2\xb5s \xc2\xb1 3.36 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n

\n

values = np.random.randint(10000, size=10000)\nids = np.random.randint(100, size=10000)\n\n%timeit Shannon(values, ids)\n1.76 ms \xc2\xb1 13.3 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit richardec(values, ids)\n29.1 ms \xc2\xb1 510 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n%timeit MadPhysicist(values, ids)\n904 \xc2\xb5s \xc2\xb1 20 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit PeptideWitch(values, ids)\n1.74 ms \xc2\xb1 16.4 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit mathfux(values, ids)\n372 \xc2\xb5s \xc2\xb1 3.54 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit bb1(values, ids)\n964 \xc2\xb5s \xc2\xb1 9.35 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n%timeit bb1_modified(values, ids)\n679 \xc2\xb5s \xc2\xb1 7.86 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1000 loops each)\n

\n

@mathfux\ 的使用建议np.maximum.reduceat是迄今为止最快的，其次是修改后的 pandas 解决方案，然后是 lexsorting。pandas 解决方案可能与reduceat，但开销要大得多。您可以看到，因此修改后的 pandas 解决方案的增量变化很小。

\n

对于大型数组（这是唯一重要的数组），似乎单独找到最大值与问题相当，而 richardec 的理解要慢得多，尽管它没有根据输入值对输入值进行排序身份证。

\n

Answer 2

np.lexsort按多列排序。然而，这不是强制性的。您可以ids先排序，然后使用以下方法选择每个分组的最大项目numpy.maximum.reduceat

def mathfux(values, ids, return_groups=False):
    argidx = np.argsort(ids) #70% time
    ids_sort, values_sort = ids[argidx], values[argidx] #4% time
    div_points = np.r_[0, np.flatnonzero(np.diff(ids_sort)) + 1] #11% time (the most part for np.flatnonzero)
    if return_groups: 
        return ids[div_points], np.maximum.reduceat(values_sort, div_points)
    else: return np.maximum.reduceat(values_sort, div_points)

mathfux(values, ids, return_groups=True)
>>> (array([0, 1, 3, 5, 6]), array([5, 3, 6, 8, 8]))
mathfux(values, ids)
>>> mathfux(values, ids)
array([5, 3, 6, 8, 8])

通常，代码的某些部分numpy可以在numba. 请注意，这np.argsort是大多数 groupby 问题的瓶颈，无法用任何其他方法替代。它不太可能很快得到改善numba或numpy。因此，您在这里达到了最佳性能，并且无法在进一步优化中做太多事情。