NumPy二维数组逐行按规则右移并补零的实现方法

完全可以实现上述效果，以下方案兼容任意行列数的输入二维数组，不会因形状差异触发报错。

核心规则对齐 #

• 每行仅提取前3个元素，超出前3位的元素直接丢弃，不参与后续计算
• 行偏移规则：以数组最底部的最后一行为基准，该行右移位数为0；向上遍历每一行，右移位数逐行加1
• 所有未被提取元素覆盖的位置，全部填充为0
• 自动适配边界：如果偏移后元素超出数组列范围，超出部分自动丢弃，不会越界报错

实现方案1：逐行索引赋值（易读易调试） #

核心逻辑很简单：先生成一个和输入数组形状完全一致的全零数组，逐行计算偏移量后，把提取出来的元素写到对应位置就行。

import numpy as np

def shift_rows(arr):
    # 生成和输入形状完全一致的全0结果数组
    result = np.zeros_like(arr)
    row_count, col_count = arr.shape
    # 从第一行（数组顶部）遍历到最后一行（数组底部）
    for row_idx in range(row_count):
        # 当前行的右移位数：最后一行偏移为0，往上逐行加1
        shift = row_count - 1 - row_idx
        # 提取当前行前3个有效元素
        extract_vals = arr[row_idx, :3]
        # 计算元素偏移后的起始、结束列位置
        start_col = shift
        end_col = shift + len(extract_vals)
        # 边界判断：如果起始列已经超出数组右边界，没有可写入的位置
        if start_col >= col_count:
            continue
        # 如果结束列超出右边界，截断写入范围适配数组宽度
        valid_end_col = min(end_col, col_count)
        valid_val_len = valid_end_col - start_col
        # 将有效元素写入结果数组对应位置
        result[row_idx, start_col:valid_end_col] = extract_vals[:valid_val_len]
    return result

测试验证 #

对题目给出的示例输入运行测试：

# 示例输入
test_input = np.array([
    [51, 52, 53, 54, 55, 56, 57],
    [41, 42, 43, 44, 45, 46, 47],
    [31, 32, 33, 34, 35, 36, 37],
    [21, 22, 23, 24, 25, 26, 27],
    [11, 12, 13, 14, 15, 16, 17]
])
output = shift_rows(test_input)
print(output)

运行输出和题目期望效果完全一致：

[[0 0 0 0 51 52 53]
 [0 0 0 41 42 43 0]
 [0 0 31 32 33 0 0]
 [0 21 22 23 0 0 0]
 [11 12 13 0 0 0 0]]

实现方案2：向量化索引（高效率适配大规模数组） #

对于数十万行的大规模数组，逐行循环的效率偏低，可以通过预计算索引网格的方式减少循环开销：

def shift_rows_vectorized(arr):
    result = np.zeros_like(arr)
    row_count, col_count = arr.shape
    # 预计算每行对应的右移偏移量
    shifts = np.arange(row_count -1, -1, -1).reshape(-1,1)
    # 预计算提取元素在原数组中的列位置（每行取前3个，对应原列索引0/1/2）
    orig_cols = np.array([0,1,2]).reshape(1,-1)
    # 计算偏移后元素在结果数组中的目标列位置
    target_cols = shifts + orig_cols
    # 过滤出没有超出数组列范围的有效索引
    valid_mask = (target_cols >=0) & (target_cols < col_count)
    # 预提取所有行需要保留的前3个元素
    extract_vals = arr[:, :3]
    # 逐行写入有效位置
    for row_idx in range(row_count):
        target_pos = target_cols[row_idx, valid_mask[row_idx]]
        valid_vals = extract_vals[row_idx, valid_mask[row_idx]]
        result[row_idx, target_pos] = valid_vals
    return result

边界适配说明 #

两种方案均支持任意形状的输入数组：

• 若数组列数较少，偏移后超出列范围的元素会自动丢弃，不会触发索引越界错误
• 若数组行数较少，偏移量不足以覆盖整行宽度时，剩余位置自动保留为0，符合补零要求

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