根据索引获取特定行数组的中值

时间:2015-02-19 01:47:17

标签: python arrays numpy

我有两个长度相同的数组,一个包含索引,另一个包含相应的值,即一个索引可以有多个值:

idx = [0,0,0,1,1,1,2,2,2,3,3,3,4,4,5,5...]
values = [1.2,3.1,3.1,3.1,3.3,1.2,3.3,4.1,5.4...]

我想返回一个数组,该数组包含唯一索引以及具有相同idx值的对象的中值。

e.g。

result = 
    [0, np.median([1.2,3.1,3.1])
     1, np.median([3.1,3.3,1.2])
     2, etc. ]

我的蛮力方法就是:

for idxi in np.arange(np.max(idx)):
    mask = (idxi == idx)
    medians = np.median(values[mask])
    result.append([idxi,medians])

不幸的是,这对于我的需求来说远非缓慢而且在任何情况下都非常难看。

2 个答案:

答案 0 :(得分:3)

如果您不介意依赖scipy,则函数scipy.ndimage.labeled_comprehension可以执行此操作。这是一个例子。

首先设置样本数据:

In [570]: import numpy as np

In [571]: idx = np.array([0,0,0,1,1,1,2,2,2,3,3,3,4,4,5,5])

In [572]: values = np.array([1.2,3.1,3.1,3.1,3.3,1.2,3.3,4.1,5.4,6,6,6.2,6,7,7.2,7.2])

获取独特的"标签"在idx。 (如果您已经知道最大值是N,并且您知道使用了从0到N的所有整数,则可以改为使用uniq = range(N+1)。)

In [573]: uniq = np.unique(idx)  # Or range(idx.max()+1)

In [574]: uniq
Out[574]: array([0, 1, 2, 3, 4, 5])

使用labeled_comprehension计算每个标记组的中位数:

In [575]: from scipy.ndimage import labeled_comprehension

In [576]: medians = labeled_comprehension(values, idx, uniq, np.median, np.float64, None)

In [577]: medians
Out[577]: array([ 3.1,  3.1,  4.1,  6. ,  6.5,  7.2])

如果您不介意依赖pandas,另一个选择是使用groupby类的pandas.DataFrame函数。

设置DataFrame:

In [609]: import pandas as pd

In [610]: df = pd.DataFrame(dict(labels=idx, values=values))

In [611]: df
Out[611]: 
    labels  values
0        0     1.2
1        0     3.1
2        0     3.1
3        1     3.1
4        1     3.3
5        1     1.2
6        2     3.3
7        2     4.1
8        2     5.4
9        3     6.0
10       3     6.0
11       3     6.2
12       4     6.0
13       4     7.0
14       5     7.2
15       5     7.2

使用groupby对数据进行分组使用labels列,然后计算群组的中位数:

In [612]: result = df.groupby('labels').median()

In [613]: result
Out[613]: 
        values
labels        
0          3.1
1          3.1
2          4.1
3          6.0
4          6.5
5          7.2

免责声明:我没有在大型阵列上尝试过这些建议,所以我不知道他们的表现与你的暴力解决方案或@Ashwini的答案相比如何。

答案 1 :(得分:2)

对于idx数组,您可以使用numpy.unique获取唯一项,并从另一个数组中获取相应的值,我们可以使用numpy.diffnumpy.where来获取项目的索引更改。使用这些索引,我们可以使用numpy.array_split拆分值数组,然后在其项目上应用np.mean

>>> idx = np.array([0,0,0,1,1,1,2,2,2,3,3,3,4,4,5,5])
>>> values = np.array([1.2,3.1,3.1,3.1,3.3,1.2,3.3,4.1,5.4] + [1]*7)
>>> indices = np.where(np.diff(idx) != 0)[0] + 1
>>> map(np.mean, np.array_split(values, indices))
[2.4666666666666668, 2.5333333333333337, 4.2666666666666666, 1.0, 1.0, 1.0]
>>> np.unique(idx)
array([0, 1, 2, 3, 4, 5])
>>> np.dstack((_, __))[0]
array([[ 0.        ,  2.46666667],
       [ 1.        ,  2.53333333],
       [ 2.        ,  4.26666667],
       [ 3.        ,  1.        ],
       [ 4.        ,  1.        ],
       [ 5.        ,  1.        ]])