Pandas 有一个广泛使用的 groupby根据相应的映射拆分 DataFrame 的工具,您可以从中对每个子组应用计算并重新组合结果。
这可以在没有原生 Python for 循环的情况下在 NumPy 中灵活地完成吗?使用 Python 循环,这看起来像:
>>> import numpy as np
>>> X = np.arange(10).reshape(5, 2)
>>> groups = np.array([0, 0, 0, 1, 1])
# Split up elements (rows) of `X` based on their element wise group
>>> np.array([X[groups==i].sum() for i in np.unique(groups)])
array([15, 30])
上面15是X
前三行的和,30是剩下两行的和。
我所说的“灵活”只是指我们不专注于某个特定的计算,例如求和、计数、最大值等,而是将任何计算传递给分组数组。
如果没有,是否有比上述方法更快的方法?
最佳答案
如何使用 scipy 稀疏矩阵
import numpy as np
from scipy import sparse
import time
x_len = 500000
g_len = 100
X = np.arange(x_len * 2).reshape(x_len, 2)
groups = np.random.randint(0, g_len, x_len)
# original
s = time.time()
a = np.array([X[groups==i].sum() for i in np.unique(groups)])
print(time.time() - s)
# using scipy sparse matrix
s = time.time()
x_sum = X.sum(axis=1)
b = np.array(sparse.coo_matrix(
(
x_sum,
(groups, np.arange(len(x_sum)))
),
shape=(g_len, x_len)
).sum(axis=1)).ravel()
print(time.time() - s)
#compare
print(np.abs((a-b)).sum())
我电脑上的结果
0.15915322303771973
0.012875080108642578
0
快 10 倍以上。
更新!
让我们对@Paul Panzer 和@Daniel F 的答案进行基准测试。它只是求和基准。
import numpy as np
from scipy import sparse
import time
# by @Daniel F
def groupby_np(X, groups, axis = 0, uf = np.add, out = None, minlength = 0, identity = None):
if minlength < groups.max() + 1:
minlength = groups.max() + 1
if identity is None:
identity = uf.identity
i = list(range(X.ndim))
del i[axis]
i = tuple(i)
n = out is None
if n:
if identity is None: # fallback to loops over 0-index for identity
assert np.all(np.in1d(np.arange(minlength), groups)), "No valid identity for unassinged groups"
s = [slice(None)] * X.ndim
for i_ in i:
s[i_] = 0
out = np.array([uf.reduce(X[tuple(s)][groups == i]) for i in range(minlength)])
else:
out = np.full((minlength,), identity, dtype = X.dtype)
uf.at(out, groups, uf.reduce(X, i))
if n:
return out
x_len = 500000
g_len = 200
X = np.arange(x_len * 2).reshape(x_len, 2)
groups = np.random.randint(0, g_len, x_len)
print("original")
s = time.time()
a = np.array([X[groups==i].sum() for i in np.unique(groups)])
print(time.time() - s)
print("use scipy coo matrix")
s = time.time()
x_sum = X.sum(axis=1)
b = np.array(sparse.coo_matrix(
(
x_sum,
(groups, np.arange(len(x_sum)))
),
shape=(g_len, x_len)
).sum(axis=1)).ravel()
print(time.time() - s)
#compare
print(np.abs((a-b)).sum())
print("use scipy csr matrix @Daniel F")
s = time.time()
x_sum = X.sum(axis=1)
c = np.array(sparse.csr_matrix(
(
x_sum,
groups,
np.arange(len(groups)+1)
),
shape=(len(groups), g_len)
).sum(axis=0)).ravel()
print(time.time() - s)
#compare
print(np.abs((a-c)).sum())
print("use bincount @Paul Panzer @Daniel F")
s = time.time()
d = np.bincount(groups, X.sum(axis=1), g_len)
print(time.time() - s)
#compare
print(np.abs((a-d)).sum())
print("use ufunc @Daniel F")
s = time.time()
e = groupby_np(X, groups)
print(time.time() - s)
#compare
print(np.abs((a-e)).sum())
标准输出
original
0.2882847785949707
use scipy coo matrix
0.012301445007324219
0
use scipy csr matrix @Daniel F
0.01046299934387207
0
use bincount @Paul Panzer @Daniel F
0.007468223571777344
0.0
use ufunc @Daniel F
0.04431319236755371
0
获胜者是 bincount 解决方案。但是csr矩阵的解法也很有意思。
关于python - 用 NumPy 向量化 groupby,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/49141969/