按 Numpy 的中位数分组(没有 Pandas)

Group by median for Numpy (without Pandas)

是否可以在不使用 pandas(并将我的数据保存在 Numpy 数组中)的情况下,根据另一列的分组来计算一列的中位数?

例如,如果这是输入:

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

我想要这个作为输出(使用第一列进行分组,然后取第二列的中位数:

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

如果您出于某种原因想避免使用 Pandas,这里有一种方法可以进行计算。请注意,在一般情况下,中位数不是整数值(除非您对其进行四舍五入或取整),因为对于 even-size 组,它将是两个最中间元素的平均值,因此您不能同时拥有单个常规数组中的整数组 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], [0, 3], [1, 4], [1, 5], [1, 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.)]