如何使用numpy有效地按值展开矩阵? [英] How to efficiently unroll a matrix by value with numpy?

问题描述

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

I have a matrix M with values 0 through N within it. I'd like to unroll this matrix to create a new matrix A where each submatrix A[i, :, :] represents whether or not M == i.

下面的解决方案使用循环.

The solution below uses a loop.

# 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操作完成此操作的方法?

Is there a faster way, or a way to do it in a single numpy operation?

推荐答案

广播比较是您的朋友:

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

 np.array_equal(A, B)
# True

想法是扩大尺寸，以使比较可以以所需的方式进行广播.

The idea is to expand the dimensions in such a way that the comparison can be broadcasted in the manner desired.

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

As pointed out by @Alex Riley in the comments, you can use np.equal.outer to avoid having to do the indexing stuff yourself,

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

np.array_equal(A, B)
# True

这篇关于如何使用numpy有效地按值展开矩阵?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！