在形状为(a,b,c)的NumPy数组元素间插零及实现逆操作

示例输入数组 #

import numpy as np

# 初始输入数组
arr = np.array([[[1., 1., 1., 1., 1.],
                 [1., 1., 1., 1., 1.],
                 [1., 1., 1., 1., 1.]],

                [[1., 1., 1., 1., 1.],
                 [1., 1., 1., 1., 1.],
                 [1., 1., 1., 1., 1.]]])
print("原数组形状:", arr.shape)  # 输出: (2, 3, 5)

一、扩展数组（插入k-1个零） #

通过索引赋值高效实现，插零数量由参数k控制（k=3时插入2个零）：

def expand_array(arr, k):
    # 构造扩展后数组的形状：最后一维长度为原长度 × k
    expanded_shape = arr.shape[:-1] + (arr.shape[-1] * k,)
    # 创建全零的目标数组
    expanded_arr = np.zeros(expanded_shape, dtype=arr.dtype)
    # 每隔k个位置填充原数组的元素
    expanded_arr[..., ::k] = arr
    return expanded_arr

# 调用扩展函数
k = 3
output = expand_array(arr, k)
print("扩展后数组形状:", output.shape)  # 输出: (2, 3, 15)
print("扩展后数组片段:\n", output[0, 0])
# 输出片段示例: [1. 0. 0. 1. 0. 0. 1. 0. 0. 1. 0. 0. 1. 0. 0.]

二、还原数组（从扩展数组恢复原数组） #

直接提取扩展数组中每隔k个位置的元素，快速还原：

def restore_array(expanded_arr, k):
    return expanded_arr[..., ::k]

# 调用还原函数
restored_arr = restore_array(output, k)
print("还原后数组形状:", restored_arr.shape)  # 输出: (2, 3, 5)
print("还原结果与原数组是否一致:", np.array_equal(restored_arr, arr))  # 输出: True

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