将单线从 Numpy Python 转换为 Julia 涉及将一个二维数组映射到另一个二维数组

Question

Julia beginner here. I need to so some jugglery with a couple of matrices. My goal is as follows:

Given a certain matrix Matrix1 as:

, and a binary matrix Matrix2 as this:

, I want to allocate the elements from Matrix1 to Matrix2 such that I have a final matrix Matrix3, which looks like:

In Python, the following one liner worked:

Matrix3= Matrix1.flatten()[(np.cumsum(Matrix2).reshape(Matrix2.shape)-1)] * Matrix2

Can anyone help me in writing a similar piece of code (preferably one liner if possible) in Julia?

Extension - I have gotten the answer for the above question from @cbk. As an extension to the question above. I was thinking of generalizing it for higher dimensional matrices. So, suppose Matrix1 has dimension (4,6,6) and the binary matrix Matrix2 has dimension (4,12,12). The allocation problem remains the same. How would then you approach it? Can someone kindly help me in it?

Answer 1

The main "gotchas" are that (1) you want to fill Matrix3 with the appropriate type of zero and (2) Julia is column-major, so you need to permute. This should do the trick though:

(Matrix3 = zeros(eltype(Matrix1),size(Matrix2)))'[Matrix2'[:]].= Matrix1'[:]

A similar option would be

(Matrix3 = zeros(eltype(Matrix1),size(Matrix2)))'[vec(Matrix2')].= vec(Matrix1')

but this was slightly less efficient than the above, as measured by @btime

edit: If Matrix2 originally contains integers instead of Booleans, you will need to convert before indexing with Matrix2 , eg:

(Matrix3 = zeros(eltype(Matrix1),size(Matrix2)))'[Bool.(Matrix2'[:])].= Matrix1'[:]