Python 2D numpy阵列组上最快的应用功能

假设我有一个2D数据矩阵，我想把一个函数应用到该矩阵中的组

例如：

对于每个唯一的索引，我想应用一些函数

f

例如，对于

index=1

组，函数

f

应用于值

0.556、0.492、0.148

（见第一列），对于

index=2

组，函数应用于值

0.043

此外：

函数必须将结果广播到输入数据的原始大小

每个列的组都是唯一的。您可以在上面的示例中看到这一点，其中每个组只包含同一列中的值

那么，在Python中执行此操作最快的方法是什么

我目前正在做以下工作（使用随机数据[2000x500]和每列5个随机组）：

---

用我的硬件，这个计算大约需要10秒才能完成。有没有更快的方法可以做到这一点？

不确定是否是最快的方法，但这种矢量化解决方案要快得多：

import numpy as np
import time

np.random.seed(0)
rows = 2000
cols = 500
ngroup = 5

data = np.random.rand(rows,cols)
groups = np.random.randint(ngroup, size=(rows,cols)) + 10*np.tile(np.arange(cols),(rows,1))

t = time.perf_counter()
# Flatten the data
dataf = data.ravel()
groupsf = groups.ravel()
# Sort by group
idx_sort = groupsf.argsort()
datafs = dataf[idx_sort]
groupsfs = groupsf[idx_sort]
# Find group bounds
idx = np.nonzero(groupsfs[1:] > groupsfs[:-1])[0]
idx = np.concatenate([[0], idx + 1, [len(datafs)]])
# Sum by groups
a = np.add.reduceat(datafs, idx[:-1])
# Count group elements
c = np.diff(idx)
# Compute group means
m = a / c
# Repeat means and counts to match data shape
means = np.repeat(m, c)
counts = np.repeat(c, c)
# Compute variance and std
v = np.add.reduceat(np.square(datafs - means), idx[:-1]) / c
s = np.sqrt(v)
# Repeat stds
stds = np.repeat(s, c)
# Compute result values
resultfs = (datafs - means) / stds
# Undo sorting
idx_unsort = np.empty_like(idx_sort)
idx_unsort[idx_sort] = np.arange(len(idx_sort))
resultf = resultfs[idx_unsort]
# Reshape back
result = np.reshape(resultf, data.shape)
print(time.perf_counter() - t)
# 0.09932469999999999

# Previous method to check result
t = time.perf_counter()
result_orig= np.zeros(data.shape)
f = lambda x: (x-np.average(x))/np.std(x)
for group in np.unique(groups):
    location = np.where(groups == group)
    group_data = data[location[0],location[1]]
    result_orig[location[0],location[1]] = f(group_data)
print(time.perf_counter() - t)
# 6.0592527

print(np.allclose(result, result_orig))
# True

编辑：要计算中间值，可以执行以下操作：

# Flatten the data
dataf = data.ravel()
groupsf = groups.ravel()
# Sort by group and value
idx_sort = np.lexsort((dataf, groupsf))
datafs = dataf[idx_sort]
groupsfs = groupsf[idx_sort]
# Find group bounds
idx = np.nonzero(groupsfs[1:] > groupsfs[:-1])[0]
idx = np.concatenate([[0], idx + 1, [len(datafs)]])
# Count group elements
c = np.diff(idx)
# Meadian index
idx_median1 = c // 2
idx_median2 = idx_median1 + (c % 2) - 1
idx_median1 += idx[:-1]
idx_median2 += idx[:-1]
# Get medians
meds = 0.5 * (datafs[idx_median1] + datafs[idx_median2])

---

这里的诀窍是使用而不仅仅是按组和值进行排序

MED

将是一个数组，其中包含每组的中位数，然后您可以使用

np。在其上重复

，如使用平均值，或任何您想要的方法。

最快的方法可能取决于具体的

f

。我们可以假设你的代码片段中的那一个就是你需要的吗？@jdehesa是的，要回答这个问题，让我们假设

f

是问题中指定的。史诗般的答案！从尝试理解您的解决方案中学到了很多。非常感谢您的时间和努力！你知道如何做一个

np.median.reduceat

？那么，以矢量化的方式找到各组的中位数吗？@Jean Paul我添加了一种计算中位数的方法，不是使用

reduceat

，而是以矢量化的方式。真不错，我没想到中位数的计算会如此简洁-非常感谢！

# Flatten the data
dataf = data.ravel()
groupsf = groups.ravel()
# Sort by group and value
idx_sort = np.lexsort((dataf, groupsf))
datafs = dataf[idx_sort]
groupsfs = groupsf[idx_sort]
# Find group bounds
idx = np.nonzero(groupsfs[1:] > groupsfs[:-1])[0]
idx = np.concatenate([[0], idx + 1, [len(datafs)]])
# Count group elements
c = np.diff(idx)
# Meadian index
idx_median1 = c // 2
idx_median2 = idx_median1 + (c % 2) - 1
idx_median1 += idx[:-1]
idx_median2 += idx[:-1]
# Get medians
meds = 0.5 * (datafs[idx_median1] + datafs[idx_median2])