torch.searchsorted

torch.searchsorted(sorted_sequence, values, *, out_int32=False, right=False, side=None, out=None, sorter=None) → Tensor

找到sorted_sequence中最内层维度的索引，使得如果将values中的相应值插入到这些索引之前，并按排序后的顺序排列时，sorted_sequence中对应最内层维度的顺序得以保留。返回一个与values大小相同的新的张量。更正式地说，返回的索引满足以下规则：

`sorted_sequence`	`right`	返回的索引符合要求
1D	假	`sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]`
1D	真	`sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]`
N-D	假	`sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]`
N-D	真	`sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]`

参数

sorted_sequence (Tensor) – N-D 或 1-D 张量，在最内层维度上包含单调递增的序列。如果提供了 sorter，则该序列不需要预先排序。
values (Tensor 或 Scalar) – 包含搜索值的 N 维张量或标量。

关键字参数

out_int32 (bool, 可选) – 指定输出数据类型。如果为 True，则输出数据类型为 torch.int32，否则为 torch.int64。默认值为 False，即默认输出数据类型为 torch.int64。
right (bool, 可选) – 如果为 False，返回找到的第一个合适的位置；如果为 True，则返回最后一个这样的索引。如果没有合适的索引被找到，对于非数值（如 nan、inf）值返回 0，或者在sorted_sequence 中的 最内层 维度大小（即超出该维度最后一个索引的位置）。换句话说，如果为 False，则获取每个 values 值在相应 最内层 维度上的下界索引；如果为 True，则获取上界索引。默认值为 False。side 与此参数功能相同，但更推荐使用它。如果将 side 设置为 "left" 而此参数设置为 True，则会引发错误。
side (str, optional) – 这与 right 参数相同，但更推荐使用。"left" 对应于 right 的 False，而 "right" 对应于 right 的 True。如果将此参数设置为 "left" 而 right 为 True，则会引发错误。默认值为 None。
out (Tensor, 可选) – 输出张量，如果提供的话必须与values大小相同。
sorter (LongTensor, optional) – 可选参数，是一个与未排序的sorted_sequence形状相同的张量，包含在最内层维度上对其进行升序排列的一系列索引。

示例：

>>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]])
>>> sorted_sequence
tensor([[ 1,  3,  5,  7,  9],
        [ 2,  4,  6,  8, 10]])
>>> values = torch.tensor([[3, 6, 9], [3, 6, 9]])
>>> values
tensor([[3, 6, 9],
        [3, 6, 9]])
>>> torch.searchsorted(sorted_sequence, values)
tensor([[1, 3, 4],
        [1, 2, 4]])
>>> torch.searchsorted(sorted_sequence, values, side='right')
tensor([[2, 3, 5],
        [1, 3, 4]])

>>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9])
>>> sorted_sequence_1d
tensor([1, 3, 5, 7, 9])
>>> torch.searchsorted(sorted_sequence_1d, values)
tensor([[1, 3, 4],
        [1, 3, 4]])