Python NumPy检查数组行列索引是否存在并按需自动扩展尺寸

NumPy原生数组为固定尺寸结构，无法直接对超出当前索引范围的坐标赋值，可通过「越界校验-按需补零扩展-赋值」的逻辑实现通用动态赋值能力，具体实现如下：

核心逻辑 #

• 传入待操作数组、目标坐标(行索引, 列索引)、待写入值
• 对比目标索引和当前数组的行列尺寸，计算满足赋值要求的最小数组尺寸
• 仅当坐标越界时重建对应尺寸的全零数组，将原数组内容拷贝到新数组的对应位置
• 完成扩展后对目标坐标赋值，返回处理后的数组

通用函数代码 #

import numpy as np

def auto_expand_assign(arr, target_row, target_col, value):
    curr_rows, curr_cols = arr.shape
    # 索引从0开始计数，目标尺寸需要比最大索引大1
    need_rows = max(curr_rows, target_row + 1)
    need_cols = max(curr_cols, target_col + 1)

    # 无越界时直接赋值，避免额外内存开销
    if need_rows == curr_rows and need_cols == curr_cols:
        arr[target_row, target_col] = value
        return arr
    
    # 构建新尺寸的全零数组，拷贝原数组内容到左上角
    expanded_arr = np.zeros((need_rows, need_cols), dtype=arr.dtype)
    expanded_arr[:curr_rows, :curr_cols] = arr
    # 写入目标值
    expanded_arr[target_row, target_col] = value
    return expanded_arr

使用示例 #

对应2×2初始矩阵给(3,1)坐标赋值1的场景，调用方式如下：

# 初始化2*2全零矩阵
matrix = np.zeros([2, 2])
# 对(3,1)位置写入1，自动触发扩展
matrix = auto_expand_assign(matrix, target_row=3, target_col=1, value=1)
print(matrix)

运行输出：

[[0. 0.]
 [0. 0.]
 [0. 0.]
 [0. 1.]]

说明 #

• 函数默认保留原数组的数据类型，不会因为扩展操作修改存储格式
• 非越界场景下不会触发数组重建，和原生NumPy赋值性能一致
• 扩展区域默认填充0，和np.zeros的初始化逻辑对齐，如果需要自定义填充值，可将代码中np.zeros替换为对应填充逻辑的初始化方法
• 支持任意合法非负整数坐标的赋值操作，无需提前预估矩阵最终尺寸

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