Numpy 间接索引

Numpy indirect indexing

我正在尝试使用另一个包含索引的数组对一个或多个数组执行聚合函数。这些索引可能包含重复项,需要根据聚合函数进行处理(我对执行此“间接索引”的一般方法感兴趣,所以我希望我不需要区分聚合函数)。

例如,假设我们想通过中的索引从v中的元素中获得和w ix.

ix = [  0,   7,   0,   1,   7,   3,   0,   2,   2, 5, 6, 4]
v =  [100, 200, 300, 400, 500, 600, 700, 800, 900, 0, 0, 0]

=>

#       0                  1               2    3  4  5  6  7
w = [1100 (100+300+700), 400, 1700 (800+900), 600, 0, 0, 0, 700 (200+500)]

sum 可能很简单,但例如加权平均值会比较棘手(v1 的乘法和 v2,然后折叠成 w)。有没有 array/numpy 方法可以做到这一点?

试试这个:

[np.sum(v[ix == [x]]) for x in range(ix.max() + 1)]

结果:

[1100  400 1700  600    0    0    0  700]

<script type="text/javascript" src="//cdn.datacamp.com/dcl-react.js.gz"></script>

<div data-datacamp-exercise data-lang="python">
  <code data-type="sample-code">
import numpy as np
ix = np.array([0, 7, 0, 1, 7, 3, 0, 2, 2, 5, 6, 4])
v = np.array([100, 200, 300, 400, 500, 600, 700, 800, 900, 0, 0, 0])

print([np.sum(v[ix == [x]]) for x in range(ix.max() + 1)])
  </code>
</div>

您正在查找 groupby 操作。 Pandas 对这种事情有相当广泛的 api 并且将 numpy 包装在引擎盖下所以你可以获得矢量化(与 numpy 一样快的一些操作)。这是一个例子:

import pandas as pd
ix = [  0,   7,   0,   1,   7,   3,   0,   2,   2, 5, 6, 4]
v =  [100, 200, 300, 400, 500, 600, 700, 800, 900, 0, 0, 0]

df = pd.DataFrame(zip(ix, v), columns=["idx", "v"])

# groupby the index, apply a sum function, convert type to numpy:
# array([1100,  400, 1700,  600,    0,    0,    0,  700])
w = df.groupby(df.idx).v.sum().to_numpy()

为了方便起见,您可以进行更复杂的计算并使用重载算术运算:

df["weights"] = np.random.rand(len(df))
df["weights"].mul(df["v"]).groupby("idx").sum()

而且性能一般:

n = 1000000
df = pd.DataFrame({"idx": np.random.choice(10, n), "v": np.random.rand(n)})

%timeit df.groupby("idx")["v"].sum()
# 11.7 ms ± 214 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

作为多功能性的展示:您可以应用更多奇特的函数,例如每个组的调和平均数(apply 有点慢):

from scipy.stats.mstats import hmean
%timeit df.groupby("idx").apply(hmean)
# 51.3 ms ± 1.74 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
idx
0    0.083368
1    0.049457
2    0.077801
3    0.074263
4    0.065142
5    0.035001
6    0.080105
7    0.002465
8    0.076336
9    0.036461

或自定义函数:

def my_func(rows):
     return np.max(rows)/np.min(rows)

%timeit df.groupby("idx")["v"].apply(my_func)
# 46.6 ms ± 2.76 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
idx
0    8.265517e+04
1    8.900603e+05
2    1.874362e+05
3    1.419228e+05
4    4.722633e+05
5    1.382114e+06
6    1.000876e+05
7    3.939510e+07
8    7.747462e+04
9    8.919914e+05

一种快速的 numpy 方法:

In [107]: ix = np.array([  0,   7,   0,   1,   7,   3,   0,   2,   2, 5, 6, 4])
     ...: v =  np.array([100, 200, 300, 400, 500, 600, 700, 800, 900, 0, 0, 0])
In [108]: 
In [108]: np.bincount(ix,v)
Out[108]: array([1100.,  400., 1700.,  600.,    0.,    0.,    0.,  700.])

另一个,不是那么快,但可能更灵活(使用其他 ufunc):

In [119]: a = np.zeros(8,int)
     ...: np.add.at(a, ix,v)
     ...: a
     ...: 
     ...: 
Out[119]: array([1100,  400, 1700,  600,    0,    0,    0,  700])

这个小例子的时间:

In [121]: timeit [np.sum(v[ix == [x]]) for x in range(ix.max() + 1)]
159 µs ± 311 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [122]: %%timeit
     ...: df = pd.DataFrame(zip(ix, v), columns=["idx", "v"])
     ...: w = df.groupby(df.idx).v.sum().to_numpy()
1.48 ms ± 884 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [123]: timeit np.bincount(ix,v)
2.15 µs ± 6.79 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [124]: %%timeit
     ...: a = np.zeros(8,int)
     ...: np.add.at(a, ix,v)
     ...: a
9.4 µs ± 348 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)