找到每个唯一bin的最大位置(binargmax)

Question

建立

假设我有

bins = np.array([0, 0, 1, 1, 2, 2, 2, 0, 1, 2])
vals = np.array([8, 7, 3, 4, 1, 2, 6, 5, 0, 9])
k = 3

我需要通过唯一bin中的最大值的位置bins.

# Bin == 0
#  ? ?           ?
# [0 0 1 1 2 2 2 0 1 2]
# [8 7 3 4 1 2 6 5 0 9]
#  ? ?           ?
#  ?
# [0 1 2 3 4 5 6 7 8 9]
# Maximum is 8 and happens at position 0

(vals * (bins == 0)).argmax()

0

---

# Bin == 1
#      ? ?         ?
# [0 0 1 1 2 2 2 0 1 2]
# [8 7 3 4 1 2 6 5 0 9]
#      ? ?         ?
#        ?
# [0 1 2 3 4 5 6 7 8 9]
# Maximum is 4 and happens at position 3

(vals * (bins == 1)).argmax()

3

---

# Bin == 2
#          ? ? ?     ?
# [0 0 1 1 2 2 2 0 1 2]
# [8 7 3 4 1 2 6 5 0 9]
#          ? ? ?     ?
#                    ?
# [0 1 2 3 4 5 6 7 8 9]
# Maximum is 9 and happens at position 9

(vals * (bins == 2)).argmax()

9

---

这些功能很苛刻,甚至不能用于负值.

题

如何使用Numpy以最有效的方式获得所有这些值？

我试过的.

def binargmax(bins, vals, k):
  out = -np.ones(k, np.int64)
  trk = np.empty(k, vals.dtype)
  trk.fill(np.nanmin(vals) - 1)

  for i in range(len(bins)):
    v = vals[i]
    b = bins[i]
    if v > trk[b]:
      trk[b] = v
      out[b] = i

  return out

binargmax(bins, vals, k)

array([0, 3, 9])

---

链接测试和验证

Answer 1

该numpy_indexed库:

我知道这不是技术上的numpy,但是numpy_indexed库有一个矢量化的group_by功能,这是完美的,只是想作为我经常使用的替代方案分享:

>>> import numpy_indexed as npi
>>> npi.group_by(bins).argmax(vals)
(array([0, 1, 2]), array([0, 3, 9], dtype=int64))

使用简单pandas groupby和idxmax:

df = pd.DataFrame({'bins': bins, 'vals': vals})
df.groupby('bins').vals.idxmax()

用一个 sparse.csr_matrix

对于非常大的输入,此选项非常快.

sparse.csr_matrix(
    (vals, bins, np.arange(vals.shape[0]+1)), (vals.shape[0], k)
).argmax(0)

# matrix([[0, 3, 9]])

性能

功能

def chris(bins, vals, k):
    return npi.group_by(bins).argmax(vals)

def chris2(df):
    return df.groupby('bins').vals.idxmax()

def chris3(bins, vals, k):
    sparse.csr_matrix((vals, bins, np.arange(vals.shape[0] + 1)), (vals.shape[0], k)).argmax(0)

def divakar(bins, vals, k):
    mx = vals.max()+1

    sidx = bins.argsort()
    sb = bins[sidx]
    sm = np.r_[sb[:-1] != sb[1:],True]

    argmax_out = np.argsort(bins*mx + vals)[sm]
    max_out = vals[argmax_out]
    return max_out, argmax_out

def divakar2(bins, vals, k):
    last_idx = np.bincount(bins).cumsum()-1
    scaled_vals = bins*(vals.max()+1) + vals
    argmax_out = np.argsort(scaled_vals)[last_idx]
    max_out = vals[argmax_out]
    return max_out, argmax_out


def user545424(bins, vals, k):
    return np.argmax(vals*(bins == np.arange(bins.max()+1)[:,np.newaxis]),axis=-1)

def user2699(bins, vals, k):
    res = []
    for v in np.unique(bins):
        idx = (bins==v)
        r = np.where(idx)[0][np.argmax(vals[idx])]
        res.append(r)
    return np.array(res)

def sacul(bins, vals, k):
    return np.lexsort((vals, bins))[np.append(np.diff(np.sort(bins)), 1).astype(bool)]

@njit
def piRSquared(bins, vals, k):
    out = -np.ones(k, np.int64)
    trk = np.empty(k, vals.dtype)
    trk.fill(np.nanmin(vals))

    for i in range(len(bins)):
        v = vals[i]
        b = bins[i]
        if v > trk[b]:
            trk[b] = v
            out[b] = i

    return out

建立

import numpy_indexed as npi
import numpy as np
import pandas as pd
from timeit import timeit
import matplotlib.pyplot as plt
from numba import njit
from scipy import sparse

res = pd.DataFrame(
       index=['chris', 'chris2', 'chris3', 'divakar', 'divakar2', 'user545424', 'user2699', 'sacul', 'piRSquared'],
       columns=[10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000],
       dtype=float
)

k = 5

for f in res.index:
    for c in res.columns:
        bins = np.random.randint(0, k, c)
        k = 5
        vals = np.random.rand(c)
        df = pd.DataFrame({'bins': bins, 'vals': vals})
        stmt = '{}(df)'.format(f) if f in {'chris2'} else '{}(bins, vals, k)'.format(f)
        setp = 'from __main__ import bins, vals, k, df, {}'.format(f)
        res.at[f, c] = timeit(stmt, setp, number=50)

ax = res.div(res.min()).T.plot(loglog=True)
ax.set_xlabel("N");
ax.set_ylabel("time (relative)");

plt.show()

结果

结果大得多k(这是广播受到重创的地方):

res = pd.DataFrame(
       index=['chris', 'chris2', 'chris3', 'divakar', 'divakar2', 'user545424', 'user2699', 'sacul', 'piRSquared'],
       columns=[10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000],
       dtype=float
)

k = 500

for f in res.index:
    for c in res.columns:
        bins = np.random.randint(0, k, c)
        vals = np.random.rand(c)
        df = pd.DataFrame({'bins': bins, 'vals': vals})
        stmt = '{}(df)'.format(f) if f in {'chris2'} else '{}(bins, vals, k)'.format(f)
        setp = 'from __main__ import bins, vals, df, k, {}'.format(f)
        res.at[f, c] = timeit(stmt, setp, number=50)

ax = res.div(res.min()).T.plot(loglog=True)
ax.set_xlabel("N");
ax.set_ylabel("time (relative)");

plt.show()

从图中可以明显看出,当组的数量很少时,广播是一个很好的技巧,但是广播的时间复杂度/存储器在较高k值时增加得太快以使其具有高性能.

Answer 2

这是通过抵消每个组数据的一种方式,以便我们可以argsort一次性使用整个数据 -

def binargmax_scale_sort(bins, vals):
    w = np.bincount(bins)
    valid_mask = w!=0
    last_idx = w[valid_mask].cumsum()-1
    scaled_vals = bins*(vals.max()+1) + vals
    #unique_bins = np.flatnonzero(valid_mask) # if needed
    return len(bins) -1 -np.argsort(scaled_vals[::-1], kind='mergesort')[last_idx]

Answer 3

好的,这是我的线性时间条目,仅使用索引和np.(max|min)inum.at.它假设垃圾箱从0上升到最大(垃圾箱).

def via_at(bins, vals):
    max_vals = np.full(bins.max()+1, -np.inf)
    np.maximum.at(max_vals, bins, vals)
    expanded = max_vals[bins]
    max_idx = np.full_like(max_vals, np.inf)
    np.minimum.at(max_idx, bins, np.where(vals == expanded, np.arange(len(bins)), np.inf))
    return max_vals, max_idx

Answer 4

这个怎么样:

>>> import numpy as np
>>> bins = np.array([0, 0, 1, 1, 2, 2, 2, 0, 1, 2])
>>> vals = np.array([8, 7, 3, 4, 1, 2, 6, 5, 0, 9])
>>> k = 3
>>> np.argmax(vals*(bins == np.arange(k)[:,np.newaxis]),axis=-1)
array([0, 3, 9])

Answer 5

如果你想要阅读,这可能不是最好的解决方案,但我认为它有效

def binargsort(bins,vals):
    s = np.lexsort((vals,bins))
    s2 = np.sort(bins)
    msk = np.roll(s2,-1) != s2
    # or use this for msk, but not noticeably better for performance:
    # msk = np.append(np.diff(np.sort(bins)),1).astype(bool)
    return s[msk]

array([0, 3, 9])

说明:

lexsortvals根据排序顺序对索引进行排序bins,然后按以下顺序排序vals:

>>> np.lexsort((vals,bins))
array([7, 1, 0, 8, 2, 3, 4, 5, 6, 9])

那么你可以通过排序bins从一个索引到下一个索引的不同来掩盖:

>>> np.sort(bins)
array([0, 0, 0, 1, 1, 1, 2, 2, 2, 2])

# Find where sorted bins end, use that as your mask on the `lexsort`
>>> np.append(np.diff(np.sort(bins)),1)
array([0, 0, 1, 0, 0, 1, 0, 0, 0, 1])

>>> np.lexsort((vals,bins))[np.append(np.diff(np.sort(bins)),1).astype(bool)]
array([0, 3, 9])

Answer 6

这是一个有趣的小问题需要解决.我的方法是vals根据中的值获取索引bins.使用where得到其中指数是点True结合argmax在丘壑这些点给出结果值.

def binargmaxA(bins, vals):
    res = []
    for v in unique(bins):
        idx = (bins==v)
        r = where(idx)[0][argmax(vals[idx])]
        res.append(r)
    return array(res)

可以unique通过使用range(k)获取可能的bin值来删除调用.这会加快速度,但随着k的大小增加,性能也会下降.

def binargmaxA2(bins, vals, k):
    res = []
    for v in range(k):
        idx = (bins==v)
        r = where(idx)[0][argmax(vals[idx])]
        res.append(r)
    return array(res)

最后一次尝试,比较每个值会大大减慢速度.此版本计算排序的值数组,而不是对每个唯一值进行比较.好吧,它实际上计算排序的索引,只在需要时获取排序值,因为这样可以避免一次将val加载到内存中.性能仍然随着垃圾箱的数量而扩展,但比之前慢得多.

def binargmaxB(bins, vals):
    idx = argsort(bins)   # Find sorted indices
    split = r_[0, where(diff(bins[idx]))[0]+1, len(bins)]  # Compute where values start in sorted array
    newmax = [argmax(vals[idx[i1:i2]]) for i1, i2 in zip(split, split[1:])]  # Find max for each value in sorted array
    return idx[newmax +split[:-1]] # Convert to indices in unsorted array

基准

这是其他答案的一些基准.

3000个元素

使用更大的数据集(bins = randint(0, 30, 3000); vals = randn(3000); k = 30;)

• Divakar的171us binargmax_scale_sort2
• 209这个答案,B版
• 281us binargmax_scale_sort来自Divakar
• 329us广播版本由user545424提供
• 399us这个答案,版本A.
• 416us使用lexsort回答sacul
• piRsquared的899us参考代码

30000个元素

还有一个更大的数据集(bins = randint(0, 30, 30000); vals = randn(30000); k = 30).令人惊讶的是,这并未改变解决方案之间的相对性能.

• 1.27ms这个答案,版本B
• Divakar的2.01ms binargmax_scale_sort2
• 2.38ms由user545424播出的版本
• 2.68ms这个答案,版本A.
• 5.71ms回答sacul,使用lexsort
• piRSquared的9.12ms参考代码

编辑我没有k随着可能的bin值的增加而改变,现在我已经修复了基准更均匀.

1000个bin值

增加唯一bin值的数量也可能会对性能产生影响.Divakar和sacul的解决方案大多不受影响,而其他解决方案则具有相当大的影响力. bins = randint(0, 1000, 30000); vals = randn(30000); k = 1000

• 1.99ms binargmax_scale_sort2来自Divakar
• 3.48ms这个答案,版本B
• 6.15ms回答sacul,使用lexsort
• piRsquared的10.6ms参考代码
• 27.2ms这个答案,版本A.
• 129毫秒广播版本由user545424

编辑包括问题中参考代码的基准,它具有惊人的竞争力,尤其是更多的箱子.