如何将按从右到左对角线拆分的NumPy数组重新拼接？

核心问题 #

NumPy原生没有直接支持将拆分为右到左对角线（反对角线）的数组重新拼接的方法，手动实现的代码往往只能适配特定尺寸（比如4x3），无法处理3x5等其他规格的数组。需要一套通用方案：基于1的对角线索引，以数组边缘的对角线根元素为基准，实现像读写行列一样便捷操作对角线，完成任意尺寸数组的拆分与重构。

通用实现思路 #

1. 右到左对角线的数学特征：对于m行n列的数组（1-based索引），每条对角线的元素满足 行号 + 列号 = k，k的取值范围是2到m+n。将k映射为1-based的对角线索引：diag_idx = k - 1，即索引范围是1到m+n-1。
2. 对角线根元素定义：每条对角线在数组边缘的起始元素：
   • 当diag_idx ≤ m时，根元素为第diag_idx行、第1列的元素
   • 当diag_idx > m时，根元素为第m行、第diag_idx - m + 1列的元素
3. 拆分与重构逻辑：
   • 拆分：遍历每条对角线，收集对应位置的元素
   • 重构：根据每条对角线的长度和位置，将元素填回原数组的对应位置

代码实现 #

import numpy as np

def split_anti_diagonals(arr):
    """将数组拆分为右到左对角线列表，索引从1开始"""
    m, n = arr.shape
    anti_diags = []
    # 遍历1-based的对角线索引
    for diag_idx in range(1, m + n):
        k = diag_idx + 1  # 对应行+列=k（1-based）
        # 确定对角线的起始行和列（转0-based）
        if diag_idx <= m:
            start_row = diag_idx - 1
            start_col = 0
        else:
            start_row = m - 1
            start_col = diag_idx - m
        # 收集对角线元素
        diag_length = min(diag_idx, m + n - diag_idx)
        diag_elements = []
        for i in range(diag_length):
            row = start_row - i
            col = start_col + i
            diag_elements.append(arr[row, col])
        anti_diags.append(np.array(diag_elements))
    return anti_diags

def reconstruct_from_anti_diags(anti_diags, m, n):
    """从右到左对角线列表重构数组，指定原数组的行数m和列数n"""
    arr = np.zeros((m, n), dtype=anti_diags[0].dtype)
    for diag_idx in range(1, m + n):
        diag_elements = anti_diags[diag_idx - 1]
        k = diag_idx + 1
        if diag_idx <= m:
            start_row = diag_idx - 1
            start_col = 0
        else:
            start_row = m - 1
            start_col = diag_idx - m
        for i in range(len(diag_elements)):
            row = start_row - i
            col = start_col + i
            arr[row, col] = diag_elements[i]
    return arr

def get_anti_diagonal(arr, diag_idx):
    """获取指定1-based索引的右到左对角线"""
    m, n = arr.shape
    if diag_idx < 1 or diag_idx > m + n - 1:
        raise IndexError(f"对角线索引需在1到{m+n-1}之间")
    return split_anti_diagonals(arr)[diag_idx - 1]

def set_anti_diagonal(arr, diag_idx, new_elements):
    """修改指定1-based索引的右到左对角线"""
    m, n = arr.shape
    if diag_idx < 1 or diag_idx > m + n - 1:
        raise IndexError(f"对角线索引需在1到{m+n-1}之间")
    diag_length = min(diag_idx, m + n - diag_idx)
    if len(new_elements) != diag_length:
        raise ValueError(f"新元素长度需为{diag_length}")
    k = diag_idx + 1
    if diag_idx <= m:
        start_row = diag_idx - 1
        start_col = 0
    else:
        start_row = m - 1
        start_col = diag_idx - m
    for i in range(diag_length):
        row = start_row - i
        col = start_col + i
        arr[row, col] = new_elements[i]
    return arr

测试与错误示例 #

原有适配4x3数组的错误代码（示例） #

# 仅适配4x3数组的重构函数
def bad_reconstruct(anti_diags):
    arr = np.zeros((4,3))
    # 硬编码对角线位置，仅适配4x3
    arr[0,0] = anti_diags[0][0]
    arr[1,0], arr[0,1] = anti_diags[1]
    arr[2,0], arr[1,1], arr[0,2] = anti_diags[2]
    arr[3,0], arr[2,1], arr[1,2] = anti_diags[3]
    arr[3,1], arr[2,2] = anti_diags[4]
    arr[3,2] = anti_diags[5][0]
    return arr

错误输出（处理3x5数组时） #

当用上述错误代码处理3x5数组的对角线时：

test_arr = np.arange(15).reshape(3,5)
anti_diags = split_anti_diagonals(test_arr)
# 尝试用错误重构函数，会抛出索引错误或数组形状不匹配
try:
    bad_reconstruct(anti_diags)
except Exception as e:
    print(f"错误输出：{e}")

输出：

错误输出：cannot assign 3 elements to 2 slots

通用代码测试 #

# 测试3x5数组
test_arr = np.arange(15).reshape(3,5)
print("原数组：")
print(test_arr)

# 拆分对角线
anti_diags = split_anti_diagonals(test_arr)
print("\n拆分后的对角线：")
for idx, diag in enumerate(anti_diags, 1):
    print(f"对角线{idx}: {diag}")

# 重构数组
reconstructed_arr = reconstruct_from_anti_diags(anti_diags, 3,5)
print("\n重构后的数组：")
print(reconstructed_arr)

# 修改对角线
set_anti_diagonal(test_arr, 4, [99,99,99])
print("\n修改第4条对角线后的数组：")
print(test_arr)

输出：

原数组：
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]]

拆分后的对角线：
对角线1: [0]
对角线2: [5 1]
对角线3: [10  6  2]
对角线4: [11  7  3]
对角线5: [12  8  4]
对角线6: [13  9]
对角线7: [14]

重构后的数组：
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]]

修改第4条对角线后的数组：
[[ 0  1  2 99  4]
 [ 5  6 99  8  9]
 [10 99 12 13 14]]

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