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	返回的索引满足
1-D	假	sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]
1-D	真	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 维或 1 维张量，在**最内层**维度上包含单调递增的序列，除非提供了 sorter，在这种情况下序列可以不必是排序的。

• values (Tensor 或 Scalar) – N 维张量或包含搜索值（们）的标量。

关键字参数

• out_int32 (bool, optional) – 指示输出数据类型。如果为 True，则为 torch.int32；否则为 torch.int64。默认值为 False，即默认输出数据类型为 torch.int64。

• right (bool, optional) – 如果为 False，则返回找到的第一个合适位置。如果为 True，则返回最后一个这样的索引。如果没有找到合适索引，则对于非数值类型（例如 nan, inf）返回 0，或者返回 sorted_sequence 中**最内层**维度的长度（比**最内层**维度的最后一个索引多一个）。换句话说，如果为 False，则获取 values 中每个值在 sorted_sequence 对应**最内层**维度上的下界索引。如果为 True，则获取上界索引。默认值为 False。 side 的功能相同且更受推荐。如果 side 设置为“left”且此参数为 True，则会报错。

• side (str, optional) – 与 right 相同但更受推荐。“left”对应于 right 的 False，“right”对应于 right 的 True。如果此参数设置为“left”而 right 为 True，则会报错。默认值为 None。

• out (Tensor, optional) – 输出张量，如果提供，则必须与 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]])