jax.numpy.lexsort#

jax.numpy.lexsort(keys, axis=-1)[source]#

按字典顺序对键序列进行排序。

JAX 对 numpy.lexsort() 的实现。

参数:
返回:

一个整数数组,形状为 keys[0].shape,给出按字典顺序排序的条目的索引。

返回类型:

Array

另请参阅

示例

具有单个键的 lexsort() 等效于 argsort()

>>> key1 = jnp.array([4, 2, 3, 2, 5])
>>> jnp.lexsort([key1])
Array([1, 3, 2, 0, 4], dtype=int32)
>>> jnp.argsort(key1)
Array([1, 3, 2, 0, 4], dtype=int32)

对于多个键,lexsort() 使用最后一个键作为主键

>>> key2 = jnp.array([2, 1, 1, 2, 2])
>>> jnp.lexsort([key1, key2])
Array([1, 2, 3, 0, 4], dtype=int32)

当打印排序后的键时,索引的含义变得更加清晰

>>> indices = jnp.lexsort([key1, key2])
>>> print(f"{key1[indices]}\n{key2[indices]}")
[2 3 2 4 5]
[1 1 2 2 2]

请注意,key2 的元素按顺序出现,并且在重复值序列中,`key1 的对应元素也按顺序出现。

对于多维输入,lexsort() 默认沿最后一个轴排序

>>> key1 = jnp.array([[2, 4, 2, 3],
...                   [3, 1, 2, 2]])
>>> key2 = jnp.array([[1, 2, 1, 3],
...                   [2, 1, 2, 1]])
>>> jnp.lexsort([key1, key2])
Array([[0, 2, 1, 3],
       [1, 3, 2, 0]], dtype=int32)

可以使用 axis 关键字选择不同的排序轴;这里我们沿前导轴排序

>>> jnp.lexsort([key1, key2], axis=0)
Array([[0, 1, 0, 1],
       [1, 0, 1, 0]], dtype=int32)