Python中将一维数组转为二维数组列表的高效实现方案

算力阿峰

2026-7-15

问题描述

我有一段总长为M*N的一维数据流，由N个通道各M长度的数据拼接而成。需要将其转换为形状为[N, K]（K≤M）的二维数组列表，以适配处理函数的输入要求。当前实现存在大量内存分配与拷贝操作，效率很低。请问有没有更优的实现方案？比如用生成器直接访问原数据切片来避免额外内存开销是否可行？或者还有其他更好的思路？

当前实现代码

测试数据生成

# 生成形状为N*M的一维测试数据
data_list: list[float] = []
for channel_no in range(N):
    data_list.extend(
        (np.arange(M) + 10 * M * channel_no).tolist()
    )

# 将一维数组重塑为二维数组，实现通道拆分
in_data: np.ndarray = np.asarray(data_list).reshape(
    N, M
)

# 生成形状为[N, K]的二维数组列表，总长度为M // K
target_data: np.ndarray = reshape_numpy_array_with_equal_blocks(
    in_array=in_data,
    target_shape=(N, K),
    drop_last=True,
    use_padding=True,
)

对应的函数实现

from typing import Tuple
import numpy as np

def reshape_numpy_array_with_equal_blocks(
    in_array: np.ndarray,
    target_shape: Tuple[int, int],
    use_padding: bool = False,
    drop_last: bool = True,
    padding_val: float = 0.0,
) -> np.ndarray:
    """将二维数组按固定块大小拆分，生成指定形状的数组列表"""

    in_array_shape = in_array.shape
    if len(in_array_shape) != 2:
        raise NotImplementedError()
    ret_array = []
    for i in range(0, in_array_shape[-1], target_shape[-1]):
        channel_array = []
        if i + target_shape[-1] <= in_array_shape[-1]:
            for channel in range(in_array_shape[0]):
                channel_array.append(
                    in_array[channel, i : i + target_shape[-1]].tolist()
                )
            ret_array.append(channel_array[:])
        else:
            if not drop_last and use_padding:
                for channel in range(in_array_shape[0]):
                    cur_data = in_array[channel, i:].tolist()
                    cur_data_len = len(cur_data)
                    cur_data.extend([padding_val] * (target_shape[-1] - cur_data_len))
                    channel_array.append(cur_data[:])
                ret_array.append(channel_array[:])
    return np.array(ret_array)

优化方案

1. 使用生成器避免内存拷贝

完全可行，生成器可以逐个产出[N, K]的切片视图（而非拷贝），全程不需要额外存储所有块的拷贝，内存开销极低。

示例实现：

def generate_blocks(in_array: np.ndarray, K: int, drop_last: bool = True, use_padding: bool = False, padding_val: float = 0.0):
    N, M = in_array.shape
    step = K
    for i in range(0, M, step):
        end = i + step
        if end > M:
            if drop_last:
                continue
            if use_padding:
                # 用numpy原生pad高效处理填充
                pad_width = (0, step - (M - i))
                block = np.pad(in_array[:, i:], ((0,0), pad_width), mode='constant', constant_values=padding_val)
                yield block
            else:
                # 直接产出剩余长度的块
                yield in_array[:, i:]
        else:
            # 产出原数组的切片视图，无拷贝
            yield in_array[:, i:end]

# 使用方式
for block in generate_blocks(in_data, K, drop_last=True, use_padding=True):
    # 处理每个[N, K]的块
    process_function(block)

注意：numpy切片是视图，只读场景下完全零额外内存，只有修改视图时才会触发写时复制。

2. 用numpy原生操作替代Python循环

如果需要一次性得到所有块的数组，用reshape+transpose组合实现批量拆分，避免Python循环的开销：

def optimized_reshape_blocks(in_array: np.ndarray, K: int, drop_last: bool = True, use_padding: bool = False, padding_val: float = 0.0):
    N, M = in_array.shape
    num_full_blocks = M // K
    remaining = M % K

    # 处理完整块：reshape+转置生成[num_blocks, N, K]的数组
    full_blocks = in_array[:, :num_full_blocks*K].reshape(N, num_full_blocks, K).transpose(1, 0, 2) if num_full_blocks > 0 else np.array([])
    
    # 处理剩余部分
    if remaining > 0 and not drop_last:
        if use_padding:
            pad_width = (0, K - remaining)
            padded_block = np.pad(in_array[:, num_full_blocks*K:], ((0,0), pad_width), mode='constant', constant_values=padding_val)
            padded_block = padded_block.reshape(1, N, K)
            return np.concatenate([full_blocks, padded_block], axis=0) if num_full_blocks >0 else padded_block
        else:
            remaining_block = in_array[:, num_full_blocks*K:].reshape(1, N, remaining)
            return np.concatenate([full_blocks, remaining_block], axis=0) if num_full_blocks >0 else remaining_block
    
    return full_blocks

此方案利用numpy的C级矢量操作，效率远高于原实现，且reshape和transpose在数组连续时都是视图操作，无内存拷贝。

3. 优化原数据生成步骤

原测试数据生成用了列表拼接+转numpy数组，可直接用numpy生成，避免列表的内存开销：

# 直接生成[N, M]的数组，再展平为一维（如果需要一维数据流）
in_data = np.arange(M) + 10*M*np.arange(N)[:, np.newaxis]
# 生成一维数据流：
data_1d = in_data.flatten()

关键优化点总结

避免频繁tolist()转换：列表与numpy数组的转换会产生大量拷贝，尽量全程用numpy数组操作。
利用numpy视图：切片、reshape、transpose等操作在数组连续时均为视图，无内存拷贝。
流式处理用生成器：如果处理函数逐个处理块，生成器可完全避免存储所有块的内存开销。
矢量操作替代循环：numpy的C实现循环远快于Python循环，大幅提升效率。

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