将2D矩阵的1s广播到1-hot 3D矩阵中

Question

给出2D矩阵A，如何广播以创建3D矩阵B？

A = [[1 1 1 1 0];[0 0 0 0 1]]
B = cat([[1 0 0 0 0];[0 0 0 0 0]],
        [[0 1 0 0 0];[0 0 0 0 0]], 
        [[0 0 1 0 0];[0 0 0 0 0]], 
        [[0 0 0 1 0];[0 0 0 0 0]], 
        [[0 0 0 0 0];[0 0 0 0 1]], dims=3)

我为这个问题标记了numpy-broadcasting，因为当我使用Julia 1.2.0时，Julia广播是基于NumPy广播建模的，因此可能会有类似的共享解决方案。

Answer 1

我不确定这是否符合您对“广播”的定义，但是一种自然的方式来获得您想要的是：

julia> cat(diagm.(getindex.(Ref(A), 1:2, :))..., dims=3)
5×5×2 Array{Int64,3}:
[:, :, 1] =
 1  0  0  0  0
 0  1  0  0  0
 0  0  1  0  0
 0  0  0  1  0
 0  0  0  0  0

[:, :, 2] =
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  1

或者如果您实际上更希望避免广播（在这种情况下，这似乎更简单）：

julia> cat((diagm(A[i, :]) for i in 1:2)..., dims=3)
5×5×2 Array{Int64,3}:
[:, :, 1] =
 1  0  0  0  0
 0  1  0  0  0
 0  0  1  0  0
 0  0  0  1  0
 0  0  0  0  0

[:, :, 2] =
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  1