python - 如何使用 numpy 有效地按值展开矩阵？

Question

我有一个矩阵 M，其中的值从 0 到 N。我想展开这个矩阵以创建一个新矩阵 A，其中每个子矩阵 A[i, :, :] 代表是否 M == i。

下面的解决方案使用了一个循环。

# Example Setup
import numpy as np

np.random.seed(0)
N = 5
M = np.random.randint(0, N, size=(5,5))

# Solution with Loop
A = np.zeros((N, M.shape[0], M.shape[1]))
for i in range(N):
    A[i, :, :] = M == i

这产生:

M
array([[4, 0, 3, 3, 3],
       [1, 3, 2, 4, 0],
       [0, 4, 2, 1, 0],
       [1, 1, 0, 1, 4],
       [3, 0, 3, 0, 2]])

M.shape
# (5, 5)


A 
array([[[0, 1, 0, 0, 0],
        [0, 0, 0, 0, 1],
        [1, 0, 0, 0, 1],
        [0, 0, 1, 0, 0],
        [0, 1, 0, 1, 0]],
       ...
       [[1, 0, 0, 0, 0],
        [0, 0, 0, 1, 0],
        [0, 1, 0, 0, 0],
        [0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0]]])

A.shape
# (5, 5, 5)

有没有更快的方法，或者在单个 numpy 操作中完成它的方法？

Answer 1

广播比较是你的 friend :

B = (M[None, :] == np.arange(N)[:, None, None]).view(np.int8)

 np.array_equal(A, B)
# True

这个想法是以这样一种方式扩展维度，以便可以以所需的方式广播比较。

---

正如@Alex Riley 在评论中指出的那样，您可以使用 np.equal.outer 来避免自己做索引工作，

B = np.equal.outer(np.arange(N), M).view(np.int8)

np.array_equal(A, B)
# True