如何将大型Numpy复数数组按特定索引规则拆分为三个子数组

我有一个大型NumPy复数数组，示例数据如下：

import numpy as np
arr = np.array([
    0.2952941 -0.22235294j, -0.4027451 +0.2090196j,
    -0.19882353+0.2717647j,  0.17764705-0.1282353j,
    0.38156864-0.23803921j,  0.25607842-0.28509805j,
    -0.4262745 +0.42078432j, -0.33215687+0.17764705j,
    -0.23019607+0.19333333j
], dtype=np.complex64)

需要将其拆分为三个NumPy数组，规则为：

• a 包含原数组第1、4、7个元素（对应索引0、3、6）
• b 包含原数组第2、5、8个元素（对应索引1、4、7）
• c 包含原数组第3、6、9个元素（对应索引2、5、8）

期望得到的结果：

a = np.array([0.2952941 -0.22235294j, 0.17764705-0.1282353j, -0.4262745 +0.42078432j])
b = np.array([-0.4027451 +0.2090196j, 0.38156864-0.23803921j, -0.33215687+0.17764705j])
c = np.array([-0.19882353+0.2717647j, 0.25607842-0.28509805j, -0.23019607+0.19333333j])

解决方案 #

方法一：切片索引（推荐） #

利用NumPy的切片语法，通过设置步长为3直接选取对应元素：

a = arr[::3]
b = arr[1::3]
c = arr[2::3]

切片语法[start:stop:step]中，省略start和stop默认覆盖整个数组，step=3表示每隔3个元素取一个，刚好匹配需求的间隔选取规则。

方法二：重塑数组后拆分 #

如果原数组长度是3的倍数，可以先将数组重塑为3行的二维数组，再提取每行作为子数组：

reshaped_arr = arr.reshape(3, -1)
a = reshaped_arr[0]
b = reshaped_arr[1]
c = reshaped_arr[2]

reshape(3, -1)会自动计算列数（这里为3），将原一维数组转为3×3的二维数组，之后直接取对应行即可得到目标子数组。

两种方法都高效且不会产生额外数据复制（切片操作返回原数组的视图），适合处理大型数组。

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