如何以对角线方式填充narray ?

Question

我想创建一个特定的数组，形状为(a, b, c, b)，如果第2个维度的索引等于最近一个维度的索引，则将其设置为1，否则设置为0。

我写了下面的代码片段，它可以工作。有没有更优雅、更有效的方法?

import numpy as np
a, b, c = 2, 3, 5
arr = np.zeros((a, b, c, b))
l = list(range(b))
arr[:, l, :, l] = 1
print(arr)

Answer 1

你可以广播一个眼睛数组[2]:

arr = np.empty((a, b, c, b))
arr[:] = np.eye(b)[None, :, None, :]

如果你想让所有的东西都在一行中，你可以平铺眼睛数组并重塑它[3]:

arr = np.tile(np.eye(b), reps=(a,c)).reshape((a, b, c, b))

计时(你的是[1]):

[1] 3.14 µs ± 46.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
[2] 1.82 µs ± 21.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
[3] 4.5 µs ± 47 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)