如何反转非方形NumPy数组的反对角线？高效实现问询

问题描述 #

需要反转NumPy数组中每条反对角线的元素顺序，且严格保持原数组的形状（例如5×3数组处理后仍为5×3），同时要支持百万级数据的高效处理。要求通过索引映射函数实现：给定目标数组的(i,j)索引和原数组形状，返回原数组中对应的元素索引，无需修改原数组。

示例输入：

import numpy as np
src_arr = np.array([
    [ 1,  2,  4],
    [ 3,  5,  7],
    [ 6,  8, 10],
    [ 9, 11, 13],
    [12, 14, 15]
])

期望输出：

[[ 1  3  6]
 [ 2  5  9]
 [ 4  8 12]
 [ 7 11 14]
 [10 13 15]]

核心思路 #

数组的每条反对角线可通过i + j = k（k为常数）标识，反转反对角线本质是将每条k对应的元素集合逆序排列。通过推导目标索引(i,j)与原数组索引的映射关系，直接计算出原数组中对应元素的位置：

1. 计算目标索引所在反对角线的k值：k = i + j
2. 确定该反对角线在原数组中的行索引范围：x_start = max(0, k - 列数 + 1)，x_end = min(行数 - 1, k)
3. 原数组对应的行索引：src_i = x_start + x_end - i
4. 原数组对应的列索引：src_j = k - src_i

实现代码 #

import numpy as np

def 反转反对角线后目标位置对应原数组索引(i, j, 数组行数, 数组列数):
    k = i + j
    x_start = max(0, k - 数组列数 + 1)
    x_end = min(数组行数 - 1, k)
    src_i = x_start + x_end - i
    src_j = k - src_i
    return (src_i, src_j)

# 测试示例
src_arr = np.array([
    [ 1,  2,  4],
    [ 3,  5,  7],
    [ 6,  8, 10],
    [ 9, 11, 13],
    [12, 14, 15]
])
rows, cols = src_arr.shape

# 生成目标数组的所有索引，批量计算原数组对应索引
i_indices, j_indices = np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij')
src_is, src_js = 反转反对角线后目标位置对应原数组索引(i_indices, j_indices, rows, cols)

# 生成结果数组
result_arr = src_arr[src_is, src_js]

print("目标输出：")
print(result_arr)

高效性说明 #

• 该实现基于纯数值计算的索引映射，无需对原数组进行任何修改或复制操作
• 利用NumPy的广播特性，可一次性生成所有目标索引对应的原数组索引，批量获取元素，时间复杂度为O(m×n)（m为行数，n为列数），完全适配百万级数据量的处理需求

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