如何在numpy中高效实现排列的复合操作

我们有作用于3维numpy数组的双射/排列，以两个由numpy数组构成的3元组形式给出。

3维numpy数组示例如下：

import numpy as np

arr = np.array([[[3, 2, 1], [5, 0, 5], [2, 0, 1]],
                [[3, 4, 5], [4, 2, 0], [0, 1, 1]],
                [[2, 0, 5], [1, 5, 1], [0, 5, 1]],
                [[4, 3, 0], [1, 3, 3], [3, 3, 3]],
                [[2, 4, 0], [2, 1, 0], [4, 4, 4]],
                [[4, 3, 2], [2, 4, 2], [5, 5, 5]]])

对应的排列示例如下：

a, b = ((np.array([5, 2, 2, 2, 1, 1, 1, 4, 4, 4, 5, 5, 0, 0, 0, 0, 0, 0, 0, 0]),
         np.array([0, 0, 0, 0, 0, 1, 2, 0, 0, 0, 2, 1, 0, 0, 0, 1, 2, 2, 2, 1]),
         np.array([2, 2, 1, 0, 0, 0, 0, 0, 1, 2, 2, 2, 0, 1, 2, 2, 2, 1, 0, 0])),
        (np.array([4, 5, 5, 5, 2, 2, 2, 1, 1, 1, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0]),
         np.array([0, 2, 1, 0, 0, 0, 0, 0, 1, 2, 0, 0, 2, 1, 0, 0, 0, 1, 2, 2]),
         np.array([2, 2, 2, 2, 2, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 2, 2, 2, 1])))

可以通过arr[a] = arr[b]的方式直接应用该映射。

复合排列的高效实现 #

核心逻辑：排列的本质是索引映射，我们可以将三维索引扁平化为一维整数构建映射表，全程使用numpy向量化操作，无Python循环，执行效率极高。

实现代码 #

def compose(p1, p2, arr_shape):
    """
    复合两个排列，应用顺序为：先执行p1，再执行p2，等价于连续执行两次arr[p[0]] = arr[p[1]]
    参数：
    p1, p2: 输入的两个排列，每个都是(a, b)格式，a、b均为三维索引三元组
    arr_shape: 目标三维numpy数组的shape属性
    返回：
    复合后的排列(c, d)，执行arr[c] = arr[d]等价于连续应用p1、p2
    """
    # 计算各维度步长，用于三维索引和一维扁平化索引的转换
    stride0 = arr_shape[1] * arr_shape[2]
    stride1 = arr_shape[2]
    
    # 把排列的三维索引全部转成一维扁平化索引
    p1_b_flat = p1[1][0] * stride0 + p1[1][1] * stride1 + p1[1][2]
    p1_a_flat = p1[0][0] * stride0 + p1[0][1] * stride1 + p1[0][2]
    p2_b_flat = p2[1][0] * stride0 + p2[1][1] * stride1 + p2[1][2]
    
    # 构建p1的索引映射表：key是目标位置，value是源位置
    map_p1 = np.arange(np.prod(arr_shape), dtype=int)
    map_p1[p1_a_flat] = p1_b_flat
    
    # 复合p2：p2的源位置需要先经过p1的映射，得到最终的原始源位置
    new_b_flat = map_p1[p2_b_flat]
    
    # 把新的源位置一维索引拆回三维格式
    new_b0 = new_b_flat // stride0
    rem = new_b_flat % stride0
    new_b1 = rem // stride1
    new_b2 = rem % stride1
    
    # 复合后的目标位置和p2的目标位置一致
    return (p2[0], (new_b0, new_b1, new_b2))

验证示例 #

# 复合两次相同的排列
c, d = compose((a, b), (a, b), arr.shape)

# 验证两种操作结果一致
arr1 = arr.copy()
arr1[a] = arr1[b]
arr1[a] = arr1[b]

arr2 = arr.copy()
arr2[c] = arr2[d]

print(np.array_equal(arr1, arr2)) # 输出True代表复合结果正确

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