一维数组提取指定索引邻域元素并双边填充补长的实现问题

针对需求：从一维NumPy数组x中，提取索引i对应的元素及其左侧l个、右侧l个元素，结果数组固定为2l+1长度；当左侧超出数组范围（i-l < 0）时补0，右侧超出（i+l ≥ len(x)）时也补0，以下提供两种简洁的实现方法：

方法一：手动计算边界拼接补0 #

直接计算有效元素的范围，再根据边界缺失情况拼接0，逻辑清晰直观：

import numpy as np

def crop_array(arr, i, l):
    arr_len = len(arr)
    # 计算理论上的起始和结束索引（左闭右开）
    start_idx = i - l
    end_idx = i + l + 1
    
    # 计算左右需要补0的数量
    left_zero_num = max(0, -start_idx)
    right_zero_num = max(0, end_idx - arr_len)
    
    # 提取数组中的有效部分
    valid_start = max(0, start_idx)
    valid_end = min(arr_len, end_idx)
    valid_elements = arr[valid_start:valid_end]
    
    # 拼接补0后的结果
    return np.concatenate([
        np.zeros(left_zero_num, dtype=arr.dtype),
        valid_elements,
        np.zeros(right_zero_num, dtype=arr.dtype)
    ])

测试示例 #

# 基础测试
i = 5
l = 2
x = np.arange(9,)
print(crop_array(x, i, l))  # 输出：array([2, 3, 4, 5, 6])

# 左侧边界测试（i-l < 0）
i = 1
l = 2
print(crop_array(x, i, l))  # 输出：array([0, 0, 0, 1, 2])

# 右侧边界测试（i+l ≥ len(x)）
i = 7
l = 2
print(crop_array(x, i, l))  # 输出：array([5, 6, 7, 8, 0])

方法二：填充后精准切片 #

利用np.pad先给数组左右各填充l个0，再通过固定长度的切片直接获取目标范围，代码更简洁：

import numpy as np

def crop_array(arr, i, l):
    # 给数组左右各填充l个0
    padded_arr = np.pad(arr, (l, l), mode='constant', constant_values=0)
    # 切片范围：从i开始，取2l+1个元素
    return padded_arr[i : i + 2*l + 1]

测试示例 #

# 基础测试
i = 5
l = 2
x = np.arange(9,)
print(crop_array(x, i, l))  # 输出：array([2, 3, 4, 5, 6])

# 左侧边界测试
i = 1
l = 2
print(crop_array(x, i, l))  # 输出：array([0, 0, 0, 1, 2])

# 右侧边界测试
i = 7
l = 2
print(crop_array(x, i, l))  # 输出：array([5, 6, 7, 8, 0])

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