处理数组:如何避免使用"for"语句

Question

我有一个名为"a"的100000000x2数组,第一列中有索引,第二列中有相关值.我需要为每个索引获取第二列中数字的中值.这就是我用for语句做的事情:

import numpy as np
b = np.zeros(1000000)
a = np.array([[1, 2],
              [1, 3],
              [2, 3],
              [2, 4],
              [2, 6],
              [1, 4],
              ...
              ...
              [1000000,6]])
for i in xrange(1000000):
    b[i]=np.median(a[np.where(a[:,0]==i),1])

对于迭代来说显然它太慢了:任何建议？谢谢

Answer 1

这被称为"分组依据"操作.Pandas(http://pandas.pydata.org/)是一个很好的工具:

import numpy as np
import pandas as pd

a = np.array([[1.0, 2.0],
              [1.0, 3.0],
              [2.0, 5.0],
              [2.0, 6.0],
              [2.0, 8.0],
              [1.0, 4.0],
              [1.0, 1.0],
              [1.0, 3.5],
              [5.0, 8.0],
              [2.0, 1.0],
              [5.0, 9.0]])

# Create the pandas DataFrame.
df = pd.DataFrame(a, columns=['index', 'value'])

# Form the groups.
grouped = df.groupby('index')

# `result` is the DataFrame containing the aggregated results.
result = grouped.aggregate(np.median)
print result

输出:

       value
index       
1        3.0
2        5.5
5        8.5

有一些方法可以直接创建包含原始数据的DataFrame,因此您不必a首先创建numpy数组.

有关Pandas中groupby操作的更多信息:http://pandas.pydata.org/pandas-docs/dev/groupby.html

Answer 2

快速的 1 行方法：

result = [np.median(a[a[:,0]==ii,1]) for ii in np.unique(a[:,0])]

我不相信您可以做很多事情来使其在不牺牲准确性的情况下进行得更快。但这是另一种尝试，如果您可以跳过排序步骤，可能会更快：

num_in_ind = np.bincount(a[:,0])
results = [np.sort(a[a[:,0]==ii,1])[num_in_ind[ii]/2] for ii in np.unique(a[:,0])]

对于小型阵列来说，后者的速度稍快一些。不确定它是否足够快。