回答此问题可获得 20 贡献值,回答如果被采纳可获得 50 分。
<p>假设我有两个numpy数组,<code>A</code>的形状<code>(d, f)</code>和形状<code>(d,)</code>的{<cd3>},其中包含{<cd5>}中的索引</p>
<pre><code>I = np.array([0, 0, 1, 0, 2, 1])
A = np.arange(12).reshape(6, 2)
</code></pre>
<p>我正在寻找一种快速的方法来进行缩减,尤其是<code>sum</code>、<code>mean</code>和{<cd8>},对所有的切片<code>A[I == i, :]</code>;一个慢的版本应该是</p>
^{pr2}$
<p>在这种情况下</p>
<pre><code>results = [[ 2.66666667, 3.66666667],
[ 7. , 8. ],
[ 8. , 9. ]])
</code></pre>
<hr/>
<p><strong>编辑</strong>:我根据Divakar的回答和之前海报(已删除)基于<code>pandas</code>的答案进行了一些计时。在</p>
<p><em>定时代码:</em></p>
<pre><code>from __future__ import division, print_function
import numpy as np, pandas as pd
from time import time
np.random.seed(0)
d = 500000
f = 500
n = 500
I = np.hstack((np.arange(n), np.random.randint(n, size=(d - n,))))
np.random.shuffle(I)
A = np.random.rand(d, f)
def reduce_naive(A, I, op="avg"):
target_dtype = (np.float if op=="avg" else A.dtype)
results = np.zeros((I.max() + 1, A.shape[1]), dtype=target_dtype)
npop = {"avg": np.mean, "sum": np.sum, "max": np.max}.get(op)
for i in np.unique(I):
results[i, :] = npop(A[I == i, :], axis=0)
return results
def reduce_reduceat(A, I, op="avg"):
sidx = I.argsort()
sI = I[sidx]
sortedA = A[sidx]
idx = np.r_[ 0, np.flatnonzero(sI[1:] != sI[:-1])+1 ]
if op == "max":
return np.maximum.reduceat(sortedA, idx, axis=0)
sums = np.add.reduceat(sortedA, idx, axis=0)
if op == "sum":
return sums
if op == "avg":
count = np.r_[idx[1:] - idx[:-1], A.shape[0] - idx[-1]]
return sums/count.astype(float)[:,None]
def reduce_bincount(A, I, op="avg"):
ids = (I[:,None] + (I.max()+1)*np.arange(A.shape[1])).ravel()
sums = np.bincount(ids, A.ravel()).reshape(A.shape[1],-1).T
if op == "sum":
return sums
if op == "avg":
return sums/np.bincount(ids).reshape(A.shape[1],-1).T
def reduce_pandas(A, I, op="avg"):
group = pd.concat([pd.DataFrame(A), pd.DataFrame(I, columns=("i",))
], axis=1
).groupby('i')
if op == "sum":
return group.sum().values
if op == "avg":
return group.mean().values
if op == "max":
return group.max().values
def reduce_hybrid(A, I, op="avg"):
sidx = I.argsort()
sI = I[sidx]
sortedA = A[sidx]
idx = np.r_[ 0, np.flatnonzero(sI[1:] != sI[:-1])+1 ]
unq_sI = sI[idx]
m = I.max()+1
N = A.shape[1]
target_dtype = (np.float if op=="avg" else A.dtype)
out = np.zeros((m,N),dtype=target_dtype)
ss_idx = np.r_[idx,A.shape[0]]
npop = {"avg": np.mean, "sum": np.sum, "max": np.max}.get(op)
for i in range(len(idx)):
out[unq_sI[i]] = npop(sortedA[ss_idx[i]:ss_idx[i+1]], axis=0)
return out
for op in ("sum", "avg", "max"):
for name, method in (("naive ", reduce_naive),
("reduceat", reduce_reduceat),
("pandas ", reduce_pandas),
("bincount", reduce_bincount),
("hybrid ", reduce_hybrid)
("numba ", reduce_numba)
):
if op == "max" and name == "bincount":
continue
# if name is not "naive":
# assert np.allclose(method(A, I, op), reduce_naive(A, I, op))
times = []
for tries in range(3):
time0 = time(); method(A, I, op)
times.<a href="https://www.cnpython.com/list/append" class="inner-link">append</a>(time() - time0);
print(name, op, "{:.2f}".format(np.min(times)))
print()
</code></pre>
<p><em>计时</em>:</p>
<pre><code>naive sum 1.10
reduceat sum 4.62
pandas sum 5.29
bincount sum 1.54
hybrid sum 0.62
numba sum 0.31
naive avg 1.12
reduceat avg 4.45
pandas avg 5.23
bincount avg 2.43
hybrid avg 0.61
numba avg 0.33
naive max 1.19
reduceat max 3.18
pandas max 5.24
hybrid max 0.72
numba max 0.34
</code></pre>
<p>(我选择了<code>d</code>和<code>n</code>作为我的用例的典型值-我在答案中添加了numba版本的代码)。在</p>