python - 根据索引获取数组特定行的中位数

Question

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

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 值的对象的唯一索引以及中值。

例如

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])

不幸的是，这对于我的需求来说太慢了，而且无论如何都很难看。

Answer 1

如果您不介意依赖 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_compressive计算每个标记组的中位数:

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 的依赖，还有另一个选择，就是使用pandas.DataFrame类的groupby函数。

设置数据框:

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的答案相比如何。