Index a matrix using two vectors separatedly repesenting its row and col

Recently, I need to solve two problems.

  1. The first is I have two vectors, flag and src. They are the same length.
    flag = [0,1,4,0,3,0]
    src = [11,9,6,7,0,7]
    flag is a vector whose element is 0 or 1 and src contains some intergers. Then I need to get the vector
    flag_ones_idx=[1,2,4] and src_selected_idx=[9,6,0]. src_selected_idx is a vector whose elements are in the positions indicated by flag_ones_idx.
    How to perform it?

  2. Then I need to set an matrix of M whose size is 10*10 and for most elements are 0s with sparse ones.
    Here I need to set M[1,9], M[4,6], M[3,0] to the value 1,1,1(the row index is each element of flag_ones_idx while the col index is each element of src_selected_idx). Any ways to do it?

flag_ones_idx = flag.nonzero().squeeze(1)
src_selected_idx = src[flag_ones_idx]

M = torch.zeros(10, 10)
M[1, 9] = 1

My bad. I am afraid that you misunderstanding my meaning for the second question. I have update the values of flag , src and the description of my second question.

Use numpy style advanced index.

>>> import torch                                    
>>> z = torch.zeros(3,4)                            
>>> z[[0, 2], [2, 3]] = 1                           
>>> z                                               
 0  0  1  0                                         
 0  0  0  0                                         
 0  0  0  1                                         
[torch.FloatTensor of size 3x4]                     
1 Like

Nice! It is exactly what I am looking for! Thank you, SimonW!