按Numpy的中位数分组(不包括Panda)



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

例如,如果这是输入:

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

我想把它作为输出(使用第一列进行分组,然后取第二列的中值:

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

如果出于某种原因想避免使用Pandas,这里有一种计算的可能性。请注意,在一般情况下,中值不是整数值(除非将其四舍五入或求底(,因为对于偶数大小的组,它将是两个中间元素的平均值,因此不能在单个正则数组中同时具有整数组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.)]

最新更新