python - Numpy: "bump"数组的高效索引数组

Question

给定一个长度 n 0 ... k-1 中的索引数组(即 A = [0, 0, 1, 2, 1, ...] )，形成新形状数组的最有效方法是什么 (n, k) B , 这样 B[i,j] = 1如果A[i] == j和 A[i] = 0不然呢？

例如，A = [0, 0, 1, 2, 1, ...] (k=3) , 我们会得到

B = [[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 0], ...]

有没有办法在没有显式 for 循环的情况下做到这一点？

Answer 1

考虑到您构建的数组的稀疏性，您可能希望使用 Scipy 的稀疏矩阵，它具有占用内存小的优点:

import numpy
from scipy import sparse

A = numpy.array([0, 0, 1, 2, 1])
k = 3
B = sparse.coo_matrix((numpy.full(len(A), 1, dtype=int), (numpy.arange(len(A)), A)), shape=(len(A), k))

(coo_matrix() 在 Scipy's documentation 中描述)。这给出了预期的结果:

>>> B.todense()
matrix([[ 1.,  0.,  0.],
        [ 1.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  1.],
        [ 0.,  1.,  0.]])

但内存占用很小(如果 k 足够大 [大于几个单位])。为了节省更多的内存，上面的 dtype 可以做得更小(取决于您的具体需要)，使用 dtype=numpy.int8 甚至 dtype= bool 。