# 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.

``````>>> 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!