NumPy 集合操作

NumPy 内置了很多用于操作集合的函数，这些函数可以非常方便用于求两个数组的并集、交集、差集以及异或等。其中，常用的集合操作函数如下表所示。

NumPy 求并集

在 NumPy 中，我们可以使用 union1d() 函数来求两个一维数组的并集。所谓的并集，指的是两个集合 “所有元素合并去重” 后的集合。

语法：

numpy.union1d(arr1, arr2)

说明：

union1d() 函数接收 2 个一维数组作为参数，然后返回一个新数组。

示例 1：计算两个数组的并集

import numpy as np

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([3, 4, 5, 6])
result = np.union1d(arr1, arr2)

print(result)

运行结果如下。

[1 2 3 4 5 6]

分析：

numpy.union1d() 会自动去重并排序（从小到大）。对于这个例子来说，numpy.union1d(arr1, arr2) 会返回 arr1 和 arr2 的并集，也就是包含两个数组中所有唯一元素的、排序好的数组。

NumPy 求交集

在 NumPy 中，我们可以使用 intersect1d() 函数来求两个一维数组的交集。所谓的并集，指的是由两个集合 “相同元素” 组成的集合。

语法：

numpy.intersect1d(
    arr1,
    arr2,
    assume_unique=False,
    return_indices=False
)

说明：

intersect1d() 函数接收以下 4 个参数。

• arr1（必选）：是一个一维数组。
• arr2（必选）：是一个一维数组。
• assume_unique（可选，默认：False）：如果为 True，则假设输入数组是唯一的（需要自行确保无重复），可以加快计算速度。
• return_indices（可选，默认：False）：如果为 True，则返回输入数组中交集元素的索引。

intersect1d() 函数会返回一个新数组。由于这是集合操作，因此新数组会去重并自动排序（从小到大）。

示例 2：计算两个数组的交集

import numpy as np

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([3, 4, 5, 6])
result = np.intersect1d(arr1, arr2)

print(result)

运行结果如下。

[3 4]

分析：

numpy.intersect1d(arr1, arr2) 会返回一个新数组，也就是 arr1 和 arr2 的交集。

示例 3：获取元素索引（return_indices=True）

import numpy as np

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([3, 4, 5, 6])
result, a_indices, b_indices = np.intersect1d(arr1, arr2, return_indices=True)

print(result)
print(a_indices)
print(b_indices)

运行结果如下。

[3 4]
[2 3]
[0 1]

分析：

如果 np.intersect1d() 设置了 return_indices=True，那么除了返回交集之外，还会返回交集元素分别在两个原数组中的索引。

NumPy 求差集

在 NumPy 中，我们可以使用 setdiff1d() 函数来求两个一维数组的差集。所谓差集，指的是 ​ “仅属于第一个集合且不属于第二个集合” 的元素集合。

语法：

numpy.setdiff1d(
    arr1,
    arr2,
    assume_unique=False
)

说明：

setdiff1d() 函数接收以下 3 个参数。

• arr1（必选）：是一个一维数组。
• arr2（必选）：是一个一维数组。
• assume_unique（可选，默认：False）：如果为 True，则假设输入数组是唯一的，可以加快计算速度。

setdiff1d() 函数会返回一个新数组。该数组的元素存在于第一个数组、但不存在于第二个数组中。

示例 4：计算两个数组的差集

import numpy as np

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([3, 4, 5, 6])
result = np.setdiff1d(arr1, arr2)

print(result)

运行结果如下。

[1 2]

NumPy 求对称差集

在 NumPy 中，我们可以使用 setxor1d() 函数来求两个一维数组的对称差集。所谓的对称差集，指的是 “只属于其中一个集合、不同时属于两者” 的元素集合。

语法：

numpy.setxor1d(
    arr1,
    arr2,
    assume_unique=False
)

说明：

setdiff1d() 函数接收以下 3 个参数。

• arr1（必选）：是一个一维数组。
• arr2（必选）：是一个一维数组。
• assume_unique（可选，默认：False）：如果为 True，则假设输入数组是唯一的，可以加快计算速度。

setxor1d() 函数会返回一个新数组，该数组包含两个数组中的不同元素。

示例 5：计算两个数组的对称差集

import numpy as np

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([3, 4, 5, 6])
result = np.setxor1d(arr1, arr2)

print(result)

运行结果如下。

[1 2 5 6]

NumPy 求唯一元素

在 NumPy 中，我们可以使用 unique() 函数来找出数组中的唯一元素。

语法：

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

说明：

unique() 函数接收以下常用的参数。

• arr（必选）：是一个数组。
• return_index（可选，默认：False）：如果为 True，则返回唯一元素的原始数组中的索引。
• return_inverse（可选，默认：False）：如果为 True，则返回原始数组中每个元素在唯一元素数组中的索引。
• return_counts（可选，默认：False）：如果为 True，则返回每个唯一元素在原始数组中出现的次数。
• axis（可选，默认：None）：指定要操作的轴。如果为 None，则输入数组会被展平。

unique() 函数会一个新数组。该数组包含了数组 arr 中所有的唯一元素，并且结果是经过排序的。

示例 6：找出数组中的唯一元素

import numpy as np

arr = np.array([1, 2, 2, 3, 4, 4, 5])
result = np.unique(arr)

print(result)

运行结果如下。

[1 2 3 4 5]

示例 7：返回唯一元素的索引

import numpy as np

arr = np.array([1, 2, 2, 3, 4, 4, 5])
elements, indices = np.unique(arr, return_index=True)

print('元素：', elements)
print('索引：', indices)

运行结果如下。

元素：[1 2 3 4 5]
索引：[0 1 3 4 6]

分析：

通过设置 return_index=True，此时除了返回唯一元素外，还返回了每个唯一元素在原始数组中首次出现的索引。

numpy.isin() 判断元素

在 NumPy 中，我们可以使用 isin() 函数来判断一个一维数组中的每个元素是否存在于第二个数组中。

语法：

numpy.isin(
    arr1,
    arr2,
    assume_unique=False,
    invert=False
)

说明：

isin() 函数接收以下 4 个参数。

• arr1（必选）：是一个一维数组。
• arr2（必选）：是一个一维数组。
• assume_unique（可选，默认：False）：如果为 True，则假设输入数组是唯一的，可以加快计算速度。
• invert（可选，默认：False）：如果为 True，则返回一个布尔数组，表示 arr1 中的哪些元素不存在于 arr2 中。

isin() 函数会返回一个布尔数组。对于布尔数组的使用语法，另请参阅：NumPy 布尔索引。

示例 8：判断元素是否存在于另一个数组中

import numpy as np

arr1 = np.array([1, 2, 3, 4, 5])
arr2 = np.array([2, 4, 6])
result = np.isin(arr1, arr2)

print(result)

运行结果如下。

[False  True False  True False]

分析：

numpy.isin(arr1, arr2) 返回了一个布尔数组，其每个元素对应于 arr1 中的元素，表示该元素是否存在于 arr2 中。