Как эффективно развернуть матрицу по значению с помощью 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)

Есть ли более быстрый способ или способ сделать это за одну операцию?

Answer 1

Вы можете использовать некоторые трансляции здесь:

P = np.arange(N)
Y = np.broadcast_to(P[:, None], M.shape)
T = np.equal(M, Y[:, None]).astype(int)

---

Альтернативное использование indices:

X, Y = np.indices(M.shape)
Z = np.equal(M, X[:, None]).astype(int)

Answer 2

Переданное сравнение - ваш друг:

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

Answer 3

Вы можете индексировать в матрицу тождественности так:

 A = np.identity(N, int)[:, M]

или около того

 A = np.identity(N, int)[M.T].T

Или использовать новую (v1.15.0) put_along_axis

A = np.zeros((N,5,5), int)
np.put_along_axis(A, M[None], 1, 0)

Обратите внимание, что если N намного больше 5, то создание идентификационной матрицы NxN может считаться расточительным.Мы можем смягчить это, используя уловки шага:

def read_only_identity(N, dtype=float):
    z = np.zeros(2*N-1, dtype)
    s, = z.strides
    z[N-1] = 1
    return np.lib.stride_tricks.as_strided(z[N-1:], (N, N), (-s, s))