如何获取numpy数组百分位对应元素索引并提取同形状数组值

需要计算Numpy数组的np.percentile，找到这些百分位值在原数组中的对应索引，再利用这些索引从另一个同形状的Numpy数组中提取对应值。

示例代码 #

import numpy as np

np.random.seed(1)

a = np.random.randint(10, size=(3,3))
b = np.random.randint(10, size=(3,3))
c = np.random.randint(10, size=(3,3))

l = np.random.randint(11,20, size=(3,3))
m = np.random.randint(11,20, size=(3,3))
n = np.random.randint(11,20, size=(3,3))

abc = np.dstack((a,b,c))
lmn = np.dstack((l,m,n))

# 查看abc数组
print("abc数组：")
print(abc)
# 输出：
# array([[[5, 9, 7],
#         [8, 2, 9],
#         [9, 4, 1]],
# 
#        [[5, 5, 7],
#         [0, 2, 0],
#         [0, 4, 6]],
# 
#        [[1, 2, 9],
#         [7, 4, 9],
#         [6, 7, 7]]])

# 查看lmn数组
print("\nlmn数组：")
print(lmn)
# 输出：
# array([[[17, 14, 11],
#         [12, 17, 14],
#         [11, 16, 13]],
# 
#        [[12, 12, 11],
#         [19, 14, 15],
#         [19, 15, 13]],
# 
#        [[14, 19, 18],
#         [19, 12, 18],
#         [18, 15, 19]]])

# 计算abc沿axis=2的50%百分位（中位数）
print("\nabc的50%百分位值：")
print(np.percentile(abc,50,axis=2))
# 输出：
# array([[7., 8., 4.],
#        [5., 0., 4.],
#        [2., 7., 7.]])

解决方案 #

要实现从lmn中提取对应abc百分位值的元素，步骤如下：

1. 对abc沿axis=2排序，获取每个位置的元素原索引排序结果
2. 确定百分位对应的索引位置（对于50%百分位，当axis=2长度为3时，对应排序后的中间位置，即索引1）
3. 使用np.take_along_axis根据索引从lmn中提取对应值

实现代码 #

# 对abc沿axis=2排序，得到每个位置的元素原索引排序
sorted_indices = np.argsort(abc, axis=2)
# 50%百分位对应排序后的中间位置（axis=2长度为3，索引为1）
percentile_pos = 1
# 调整索引维度，适配take_along_axis的要求
target_indices = sorted_indices[:, :, percentile_pos][..., np.newaxis]
# 从lmn中提取对应值
result = np.take_along_axis(lmn, target_indices, axis=2).squeeze()

print("\n提取结果：")
print(result)
# 输出：
# array([[11, 12, 16],
#        [12, 15, 15],
#        [19, 19, 15]])

说明 #

• np.argsort(abc, axis=2)返回沿axis=2排序后的元素原索引，例如abc[0,0] = [5,9,7]排序后原索引为[0,2,1]（对应值5,7,9）
• 50%百分位对应排序后的第1个索引（0-based），即中位数所在的位置
• np.take_along_axis用于根据指定轴的索引提取元素，squeeze()用于去除多余的单维度

内容的提问来源于stack exchange，提问作者Pedro VCO