import torch


def invert_s2o(s2o):
    """
    given you've done:

        sorted, s2o = sometensor.sort()

    ... so that s2o is the index in the original tensor of position s in the sorted tensor,
    then you can call:

        o2s = invert_s2o(s2o)

    ... to get o2s, which is the index in the sorted tensor of index o in the original tensor

    assumes that s and o are 1-dimensional
    """
    assert len(s2o.size()) == 1
    N = s2o.size(0)
    sequential = torch.ones(N, device=s2o.device, dtype=torch.int64).cumsum(dim=-1) - 1
    o2s = torch.zeros_like(sequential)
    o2s.scatter_(0, s2o, sequential)
    return o2s