按Numpy的中位数分组不含熊猫

编程问答 2022-10-27

问题描述

是否可以在不使用熊猫的情况下（根据我的数据保存在Numpy数组中）基于另一列的分组来计算一列的中位数？

例如，如果这是输入：

arr = np.array([[0,1],[0,2],3],[1,4],5],6]])

我希望将此作为输出（使用第一列进行分组，然后取第二列的中位数：

ans = np.array([[0,5]])

解决方法

如果出于某种原因要避免使用熊猫，这是一种进行计算的可能性。请注意，在通常情况下，中位数不是整数值（除非将其四舍五入或取整），因为对于偶数大小的组，它是两个中间元素的平均值，因此不能同时拥有两个整数组单个规则数组中的id和中值（尽管可以在结构化数组中）。

import numpy as np

def grouped_median(group,value):
    # Sort by group and value
    s = np.lexsort([value,group])
    arr2 = arr[s]
    group2 = group[s]
    value2 = value[s]
    # Look for group boundaries
    w = np.flatnonzero(np.diff(group2,prepend=group2[0] - 1,append=group2[-1] + 1))
    # Size of each group
    wd = np.diff(w)
    # Mid points of each group
    m1 = w[:-1] + wd // 2
    m2 = m1 - 1 + (wd % 2)
    # Group id
    group_res = group2[m1]
    # Group median value
    value_res = (value2[m1] + value2[m2]) / 2  # Use `// 2` or round for int result
    return group_res,value_res

# Test
arr = np.array([[0,1],[0,2],3],[1,4],5],6]])
group_res,value_res = grouped_median(arr[:,0],arr[:,1])
# Print
for g,v in zip(group_res,value_res):
    print(g,v)
    # 0 2.0
    # 1 5.0
# As a structured array
res = np.empty(group_res.shape,dtype=[('group',group_res.dtype),('median',value_res.dtype)])
res['group'] = group_res
res['median'] = value_res
print(res)
# [(0,2.) (1,5.)]
group-bymediannumpypython