基于NumPy的3D数组内2D数组指定列无循环降序排名替换

针对形状为 (3, 5, 5) 的3D数组，要将每个内部2D数组的第4列（索引3）替换为降序排名（最高值对应行数5，最低值对应1），可以用以下两种原生NumPy向量化方案，完全避免循环：

步骤说明与代码实现 #

1. 构造示例数组 #

先修正原数组的格式错误，创建符合要求的3D数组：

import numpy as np

arr = np.array([
    [[1. , 0. , 0. , 10., 0.],
     [2. , 0. , 0. , 9. , 0.],
     [3. , 0. , 0. , 8. , 0.],
     [4. , 0. , 0. , 7. , 0.],
     [5. , 0. , 0. , 6. , 0.]],
    [[1. , 0. , 0. , 199., 0.],
     [2. , 0. , 0. , 198., 0.],
     [3. , 0. , 0. , 196., 0.],
     [4. , 0. , 0. , 190., 0.],
     [5. , 0. , 0. , 160., 0.]],
    [[1. , 0. , 0. , 999., 0.],
     [2. , 0. , 0. , 870., 0.],
     [3. , 0. , 0. , 270., 0.],
     [4. , 0. , 0. , 100., 0.],
     [5. , 0. , 0. , 80. , 0.]]
])

2. 方案一：简洁的两次argsort写法 #

利用两次argsort快速计算降序排名，是最简洁的实现方式：

# 提取所有2D数组的第4列，形状为(3,5)
target_col = arr[:, :, 3]

# 计算从1开始的降序排名
# 取反后argsort等价于原数组降序排序，第二次argsort转换为排名，加1得到1起始的排名
ranks = (-target_col).argsort(axis=1).argsort(axis=1) + 1

# 将排名替换回原数组的目标列
arr[:, :, 3] = ranks

3. 方案二：直观的索引赋值写法 #

如果需要更直观的逻辑，可通过降序索引直接赋值排名：

# 提取目标列
target_col = arr[:, :, 3]
row_count = arr.shape[1]  # 每个2D数组的行数：5

# 获取每个子数组目标列的降序索引
sorted_indices = target_col.argsort(axis=1)[:, ::-1]

# 创建空的排名数组，按降序索引依次赋值1~5
ranks = np.zeros_like(target_col)
ranks[np.arange(ranks.shape[0])[:, None], sorted_indices] = np.arange(1, row_count + 1)

# 替换回原数组
arr[:, :, 3] = ranks

验证结果 #

执行上述代码后，原数组的第4列会被替换为降序排名，最终输出为：

[[[ 1.  0.  0.  5.  0.]
  [ 2.  0.  0.  4.  0.]
  [ 3.  0.  0.  3.  0.]
  [ 4.  0.  0.  2.  0.]
  [ 5.  0.  0.  1.  0.]]

 [[ 1.  0.  0.  5.  0.]
  [ 2.  0.  0.  4.  0.]
  [ 3.  0.  0.  3.  0.]
  [ 4.  0.  0.  2.  0.]
  [ 5.  0.  0.  1.  0.]]

 [[ 1.  0.  0.  5.  0.]
  [ 2.  0.  0.  4.  0.]
  [ 3.  0.  0.  3.  0.]
  [ 4.  0.  0.  2.  0.]
  [ 5.  0.  0.  1.  0.]]]

方案优势 #

两种方案都是完全向量化的NumPy操作，没有显式循环，效率远高于遍历实现，适合处理任意规模的同结构3D数组。

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