python - 将numpy索引数组扩展为矩阵

Question

我有一个 N 整数的 numpy 数组，范围从 0 到 M （含）。我希望将它们视为 NxM 矩阵的索引，该矩阵在数组指示的每个位置都包含 1，在其他任何位置都包含 0。例如，如果给定 N=4, M=2 我有以下数组

[1, 0, 2, 1]

我想得到这个矩阵

[0 1 0]
[1 0 0]
[0 0 1]
[0 1 0]

即第 0 行在第 1 列中有 1，第 1 行在第 0 列中有 1，依此类推。

如何在 numpy 中进行这种转换？

Answer 1

这需要多维数组索引。

a = np.array([1, 0, 2, 1])
z = np.zeros(12, dtype=int).reshape(4,3)
z[np.arange(a.size), a] = 1