PyTorch paste values into tensor by row index with increasing column index
I have a tensor output into which I want to put some values. I know the row that each value should go in, but I don't have an index tensor describing the columns. Instead, if there are k values that belong to one row, they should go in columns 0, 1, ..., k-1. This is perhaps better explained with an example:

import torch

n = 4
max_cols = 5
output = torch.zeros(n, max_cols)
row_idx = torch.tensor([0, 0, 0, 0, 1, 1, 2, 2, 2, 3])
values = torch.arange(len(row_idx)).float() # the values could be anything, not just arange
# output[??] = values


Here the first 4 values should be in output at [0, 0], ... [0, 3], the next 2 values at [1, 0], [1, 1], and so on.

Here's how I'm doing this now

_, counts = torch.unique(row_idx, return_counts=True)
range_ = torch.arange(max_cols)
col_idx = torch.cat([range_[:c] for c in counts])
output[row_idx, col_idx] = values
output


tensor([[0., 1., 2., 3., 0.],
        [4., 5., 0., 0., 0.],
        [6., 7., 8., 0., 0.],
        [9., 0., 0., 0., 0.]])


Is there any more efficient way to paste these values into the appropriate positions?

(feel free to suggest a better title for this if you can think of one)

I think your solution has linear time complexity. So, I am not sure if it can be further improved. However, I think the solution you provided is not correct. Let me give an example.

For the following input:

row_idx = torch.tensor([0, 0, 1, 0, 0, 1, 2, 2, 2, 3])


Your solution outputs the following.

tensor([[4., 1., 3., 3., 0.],
        [2., 5., 2., 0., 0.],
        [6., 7., 8., 0., 0.],
        [9., 0., 0., 0., 0.]])


However, I think your expected output is:

tensor([[0., 1., 3., 4., 0.],
        [2., 5., 0., 0., 0.],
        [6., 7., 8., 0., 0.],
        [9., 0., 0., 0., 0.]])


So, I suggest the following solution that I believe is correct.

def helper(a):
    idx = a.cumsum(-1)
    id_arr = torch.ones(idx[-1], dtype=int)
    id_arr[0] = 0
    id_arr[idx[:-1]] = -a[:-1]+1
    return id_arr.cumsum(-1)

n = 4
max_cols = 5
output = torch.zeros(n, max_cols)
row_idx = torch.tensor([0, 0, 1, 0, 0, 1, 2, 2, 2, 3])
values = torch.arange(len(row_idx)).float()

count = torch.unique(row_idx, return_counts=True)[1]
col_idx = helper(count)[row_idx.argsort().argsort()]
output[row_idx, col_idx] = values
print(output)




Update

You can simply add one line to your code as follows to make it work correctly.

_, counts = torch.unique(row_idx, return_counts=True)
range_ = torch.arange(max_cols)
col_idx = torch.cat([range_[:c] for c in counts])
col_idx = col_idx[row_idx.argsort().argsort()] # &lt;== UPDATE
output[row_idx, col_idx] = values
print(output)

