numpy.unique 及其参数含义

numpy.unique(arr, return_index=False, return_inverse=False, return_counts=False, axis=None)

对原始 array 去重，且返回排序好的 array。

因为涉及到排序，axis 参数的用法有点 confusing，稍后介绍。不指定 axis 参数，将会对 arr 中所有元素进行操作，相当于把原始 array 展平（如果是多维的话）。

我们先结合这几个参数，来看几个例子。

import numpy as np

names = np.array(['Bob', 'Joe', 'Will', 'Bob', 'Will', 'Joe', 'Joe'])
np.unique(names)
"""
array(['Bob', 'Joe', 'Will'], dtype='
1
2
3
4
5
6
7

这与下面的表达是等价的：

sorted(set(names))
"""
['Bob', 'Joe', 'Will']
"""
1
2
3
4

---

我们使用 return_index 参数：

ints = np.array([3, 3, 3, 2, 2, 1, 1, 4, 4])
uints, indices = np.unique(ints, return_index=True)
1
2

uints
"""
array([1, 2, 3, 4])
"""

indices
"""
array([5, 3, 0, 7], dtype=int64)
"""
1
2
3
4
5
6
7
8
9

原始 array ints 中下标为 [5, 3, 0, 7] 的位置对应的元素确实为 [1, 2, 3, 4]。注意，这里选择的是重复元素中最左边的元素。

---

再来看 return_inverse 参数：

uints, indices = np.unique(ints, return_inverse=True)

uints
"""
array([1, 2, 3, 4])
"""

indices
"""
array([2, 2, 2, 1, 1, 0, 0, 3, 3], dtype=int64)
"""
1
2
3
4
5
6
7
8
9
10
11

根据 indices 提供的下标我们在 uints 中取值，这里用到了 NumPy 的 fancy indexing：

uints[indices]
"""
array([3, 3, 3, 2, 2, 1, 1, 4, 4])
"""
1
2
3
4

---

return_counts 参数：

uints, counts = np.unique(ints, return_counts=True)

uints
"""
array([1, 2, 3, 4])
"""

counts
"""
array([2, 2, 3, 2], dtype=int64)
"""
1
2
3
4
5
6
7
8
9
10
11

---

最后我们来看 axis 参数。我们可以直观的这样理解，对于多维数组，axis=0 会返回 unique 的行，而 axis=1会返回 unique 的列：

np.unique(np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]]), axis=0)
"""
array([[1, 0, 0],
       [2, 0, 0]])
"""
1
2
3
4
5

np.unique(np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]]), axis=1)
"""
array([[0, 1],
       [0, 1],
       [0, 2]])
"""
1
2
3
4
5
6

注意到在 axis=1 时，对去重结果排序后，列的顺序与原来的顺序已经不一样了。下面再看两个类似的例子：

a = np.array([[1, 2, 6], [1, 0, 5], [2, 3, 4]])
a
"""
array([[1, 2, 6],
       [1, 0, 5],
       [2, 3, 4]])
"""

np.unique(a, axis=1)
"""
array([[1, 2, 6],
       [1, 0, 5],
       [2, 3, 4]])
"""
1
2
3
4
5
6
7
8
9
10
11
12
13
14

np.unique(a, axis=0)
"""
array([[1, 0, 5],
       [1, 2, 6],
       [2, 3, 4]])
"""
1
2
3
4
5
6

我们注意到，在下面 axis=0 的情况中，返回的结果中前两行位置互换了。这是最后排序的结果，使用的是 lexsort（可看这篇文章《NumPy 中的排序方法》）。前两行的第一个元素都为 1，而第二个元素 0 < 2，因此第二行会排到第一行前面。

axis=1时，因为每列的第一个元素已经是有序的了，所以结果并未改变。