主页 / coding / 问题 / 78748429
Accepted
sams-studio
Asked: 2024-07-15 14:58:46 +0800 CST

在 numpy 中的多个轴上应用排列数组

  • 772

假设我有一个排列数组perm，如下所示：

perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])

如果我想将其应用于一个轴，我可以写类似如下的内容：

v = np.arange(9).reshape(3, 3)
print(v[perm])

输出：

array([[[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]],

       [[3, 4, 5],
        [6, 7, 8],
        [0, 1, 2]],

       [[0, 1, 2],
        [6, 7, 8],
        [3, 4, 5]],

       [[6, 7, 8],
        [3, 4, 5],
        [0, 1, 2]]])

现在我想同时将其应用于两个轴。我发现我可以通过以下方式实现：

np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

但我发现它效率很低，因为它必须创建一个网格，而且还需要一个 for 循环。我在这个例子中创建了一个小数组，但实际上我有很多更大的数组和很多排列，所以我真的很想有一个像单轴版本一样快速和简单的版本。

python

3 个回答

  1. Best Answer

    怎么样：

    p1 = perm[:, :, np.newaxis]
p2 = perm[:, np.newaxis, :]
v[p1, p2]

    p1和的第零轴p2只是 的“批量”维度perm，它允许您在一次操作中执行多次排列。

    的另一个维度perm与索引相对应，沿中的第一个轴p1和中的第二个轴对齐p2。由于轴是正交的，因此数组被广播，基本上就像您使用的数组一样meshgrid- 但这些仍然具有批处理维度。

    这是我通过手机能做到的最好的事情：）如果需要的话我可以稍后尝试澄清，但关键思想是广播。

    比较：

    import numpy as np
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])
v = np.arange(9).reshape(3, 3)

ref = np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

p1 = perm[:, :, np.newaxis]
p2 = perm[:, np.newaxis, :]
res = v[p1, p2]

np.testing.assert_equal(res, ref)
# passes

%timeit np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])
# 107 µs ± 20.6 µs per loop

%timeit v[perm[:, :, np.newaxis], perm[:, np.newaxis, :]]
# 3.73 µs ± 1.07 µs per loop

    广播索引的一个更简单（没有批量维度）的示例：

    import numpy as np
i = np.arange(3)
ref = np.meshgrid(i, i, indexing="ij")
res = np.broadcast_arrays(i[:, np.newaxis], i[np.newaxis, :])
np.testing.assert_equal(res, ref)
# passes

    在顶部的解决方案代码中，广播是隐式的。我们不需要调用，broadcast_arrays因为它会在索引过程中自动发生。

    • 4

  2. meshgrid你可以摆脱

    a = np.array([v[p][:,p] for p in perm])
b = np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])
print(np.all(b == a)) # True

    这比示例数组快 5 倍：

    import timeit
%timeit np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm]) # 42.7 µs
%timeit np.array([v[p][:,p] for p in perm]) # 8.18 µs

    我认为 for 循环基本上是不相关的，如果您关心进一步优化，请指定您正在使用的形状......

    • 2

  3. 要同时将排列应用于两个轴而不使用循环或网格，可以使用高级索引和广播。这是一种更有效的方法：

    import numpy as np

# Define permutations and initial array
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])
v = np.arange(9).reshape(3, 3)

# Apply permutations using broadcasting
permuted_indices = perm[:, :, None] * 3 + perm[:, None, :]
result = v.ravel()[permuted_indices]

print(result)
    1. perm[:, :, None] * 3 + perm[:, None, :]使用广播计算扁平数组中排列值的 1D 索引。
    2. v.ravel()将数组v展平。
    3. 结果是对两个轴都应用了排列的所需输出数组。

    这种方法更加高效并且避免了显式循环和meshgrid。

    • 0

相关问题