如何用Numpy高效实现数组的反锯齿化（Unzigzag）？

问题背景 #

我们有一个一维数组，它是通过如下锯齿路径采样二维矩阵得到的：

↓-------<------S  <- start
>------->------↓
↓-------<------<
>------->------E  <- end

例如，原矩阵：

B =    1 2  3  4
       5 6  7  8
       9 10 11 12

被编码为数组 A = [4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9]，矩阵行数可为奇数或偶数。

现有循环实现的还原函数效率低下，需要用Numpy实现矢量化的高效反锯齿化，同时支持多维输入：

• 输入形状 (12,) → 输出形状 (3, 4)
• 输入形状 (12, 100) → 输出形状 (3, 4, 100)
• 输入形状 (n, i, j, k, ...) → 输出形状 (n/numcols, numcols, i, j, k, ...)

矢量化解决方案 #

核心思路 #

观察采样规律：编码后的数组按行分组，每组对应原矩阵的一行。其中偶数索引的行（0-based）是原行的逆序，奇数索引的行与原行顺序一致。基于此，我们可以通过「重塑数组+选择性反转行」实现完全矢量化的高效还原。

实现代码 #

import numpy as np

def unzigzag(z, numcols):
    z = np.asarray(z)
    # 计算行数，确保输入第一维度长度能被numcols整除
    numrows = z.shape[0] // numcols
    # 重塑形状：将第一维度拆分为(numrows, numcols)，保留后续所有维度
    reshaped = z.reshape(numrows, numcols, *z.shape[1:])
    # 筛选需要反转的行：偶数索引行（0,2,4...）
    reverse_mask = np.arange(numrows) % 2 == 0
    # 对目标行沿列轴反转
    reshaped[reverse_mask] = reshaped[reverse_mask, ::-1]
    return reshaped

测试验证 #

一维输入测试

A = [4, 3, 2, 1, 5, 6, 7, 8, 12, 11, 10, 9]
B = unzigzag(A, 4)
print(B)
# 输出：
# [[ 1  2  3  4]
#  [ 5  6  7  8]
#  [ 9 10 11 12]]

多维输入测试

A_multi = np.array([[4,4], [3,3], [2,2], [1,1], [5,5], [6,6], [7,7], [8,8], [12,12], [11,11], [10,10], [9,9]])
B_multi = unzigzag(A_multi, 4)
print(B_multi.shape)  # 输出：(3, 4, 2)
print(B_multi)
# 输出：
# [[[ 1  1]
#   [ 2  2]
#   [ 3  3]
#   [ 4  4]]
# 
#  [[ 5  5]
#   [ 6  6]
#   [ 7  7]
#   [ 8  8]]
# 
#  [[ 9  9]
#   [10 10]
#   [11 11]
#   [12 12]]]

优势说明 #

该实现完全基于Numpy矢量化操作，避免了Python循环的性能开销，处理大数组时效率远高于原循环版本。同时自动适配多维输入，无需额外修改逻辑。

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