NumPy 排序、条件刷选函数


NumPy 排序、条件刷选函数

排序

NumPy 提供了多种排序方法,同时可以指定对哪个轴进行排序。常用的排序方法有:

  • sort()
  • argsort()
  • lexsort()
  • argpartition()
  • partition()

sort()

sort() 方法用于对数组进行排序,该方法会返回一个已排序的数组副本。

import numpy as np

arr = np.array([3, 1, 2])
print(np.sort(arr))  # [1 2 3]

sort() 方法默认对全部元素进行排序,如果想只对部分进行排序,可以利用切片进行处理:

arr = np.array([3, 1, 2])
print(np.sort(arr[0:2]))  # [1 3]

sort() 方法还可以对多维数组进行排序,指定 axis 参数即可:

arr = np.array([[2, 4], [1, 3]])
print(np.sort(arr, axis=1))  # [[2 4] [1 3]]

argsort()

argsort() 方法用于返回数组排序后的下标,也就是返回排序后数组中元素在原始数组中的位置。

arr = np.array([3, 1, 2])
print(np.argsort(arr))  # [1 2 0]

lexsort()

lexsort() 方法用于执行间接排序,该方法将多个序列一起排序,因此要求每个序列的元素个数相同。

a = np.array([1, 3, 2])
b = np.array([7, 6, 5])
c = np.array([9, 8, 0])
d = np.lexsort((a, b, c))
print(list(zip(a[d], b[d], c[d])))  # [(2, 5, 0), (3, 6, 8), (1, 7, 9)]

argpartition()

argpartition() 方法用于执行快速排序中的分区(即快排中的 pivot 操作)。

arr = np.array([3, 1, 2])
print(np.argpartition(arr, 1))  # [1 2 0]

partition()

partition() 方法用于实现部分排序,也就是说可以让数组在分区中有序,在分区之外无序。

arr = np.array([3, 1, 2])
print(np.partition(arr, 1))  # [1 2 3]

条件刷选函数

NumPy 的条件刷选函数可以按照指定条件选取符合条件的数组元素。

  • nonzero()
  • where()
  • extract()

nonzero()

nonzero() 方法用于返回非零元素的下标。

arr = np.array([3, 0, 2, 0, 1])
print(np.nonzero(arr))  # [(array([0, 2, 4], dtype=int64),)]

where()

where() 方法用于按照某个条件选取数组中符合条件的元素,如果没有条件,就返回所有非零元素的下标。

arr = np.array([3, 0, 2, 0, 1])
print(np.where(arr > 1))  # (array([0, 2], dtype=int64),)

extract()

extract() 方法用于从数组中提取符合条件的元素。

arr = np.array([3, 0, 2, 0, 1])
condition = arr > 1
print(np.extract(condition, arr))  # [3 2]

总结

NumPy 提供了很多种排序和条件刷选函数,可以满足我们对数组的排序和选取需求。在使用时需要根据具体情况选择适当的函数,并注意每个函数的参数和返回值,以免出现错误。