NumPy高效切片：无循环处理变长起始/结束索引的技巧

import numpy as np

# 为了结果可复现，设置随机种子
np.random.seed(42)

# 原始一维数组
a = np.random.randn(10)
print(f"原始数组 a: {a}\n")

# 定义多个切片的起始和结束索引
starts = np.array([1, 2])
ends = np.array([3, 4])
# 验证每个切片的长度M
M = ends[0] - starts[0]
print(f"切片起始索引: {starts}")
print(f"切片结束索引: {ends}")
print(f"每个切片的固定长度 M: {M}\n")

# 使用 for 循环进行切片
all_slices_for_loop = []
for s, e in zip(starts, ends):
  all_slices_for_loop.append(a[s:e])

all_slices_for_loop = np.stack(all_slices_for_loop, axis=0)
print(f"通过 for 循环得到的切片结果 (形状 {all_slices_for_loop.shape}):\n{all_slices_for_loop}")

原始数组 a: [ 0.49671415 -0.1382643  1.64768854  0.76743472 -1.28020898  0.03602528
 -0.21798363  0.50275798 -0.36239599  0.86518764]

切片起始索引: [1 2]
切片结束索引: [3 4]
每个切片的固定长度 M: 2

通过 for 循环得到的切片结果 (形状 (2, 2)):
[[-0.1382643   1.64768854]
 [ 1.64768854  0.76743472]]

# 沿用之前的 a, starts, ends, M

# 方法一：利用广播机制生成索引数组
# starts[:, None] 将 starts 从 (N,) 变为 (N, 1)
# np.arange(M) 生成 [0, 1, ..., M-1]
# 两者相加利用广播机制，生成 (N, M) 的索引数组
idx_broadcast = starts[:, None] + np.arange(M)
print(f"方法一生成的索引数组 idx_broadcast (形状 {idx_broadcast.shape}):\n{idx_broadcast}\n")

# 使用高级索引提取数据
result_broadcast = a[idx_broadcast]
print(f"方法一得到的切片结果 (形状 {result_broadcast.shape}):\n{result_broadcast}")

# 验证结果与 for 循环一致
assert np.array_equal(result_broadcast, all_slices_for_loop)
print("\n方法一结果与 for 循环结果一致。")

方法一生成的索引数组 idx_broadcast (形状 (2, 2)):
[[1 2]
 [2 3]]

方法一得到的切片结果 (形状 (2, 2)):
[[-0.1382643   1.64768854]
 [ 1.64768854  0.76743472]]

方法一结果与 for 循环结果一致。

# 沿用之前的 a, starts, ends, M

# 方法二：利用 np.linspace 生成索引数组
# np.linspace(starts, ends, num=M, dtype=int, endpoint=False)
# 这里的 starts 和 ends 都是 (N,) 数组，linspace 会对每对 (starts[i], ends[i]) 生成 M 个点
# 最终生成的是一个 (M, N) 的数组，需要转置为 (N, M)
idx_linspace = np.linspace(starts, ends, num=M, dtype=int, endpoint=False).T
print(f"方法二生成的索引数组 idx_linspace (形状 {idx_linspace.shape}):\n{idx_linspace}\n")

# 使用高级索引提取数据
result_linspace = a[idx_linspace]
print(f"方法二得到的切片结果 (形状 {result_linspace.shape}):\n{result_linspace}")

# 验证结果与 for 循环一致
assert np.array_equal(result_linspace, all_slices_for_loop)
print("\n方法二结果与 for 循环结果一致。")

方法二生成的索引数组 idx_linspace (形状 (2, 2)):
[[1 2]
 [2 3]]

方法二得到的切片结果 (形状 (2, 2)):
[[-0.1382643   1.64768854]
 [ 1.64768854  0.76743472]]

方法二结果与 for 循环结果一致。