NumPy数组“缘分”大揭秘：告别循环，用向量化秒查对应关系！

引言：当代码遇上“缘分”

PYTHON
# 朴素的循环方法
common = []
for item in array_a:
    if item in array_b:
        common.append(item)

什么是向量化？

• 循环（Looping）：就像你一个一个地拿起苹果，走过去放下，再回来拿下一个。重复劳动，累且慢。
• 向量化（Vectorization）：就像你直接抱起整个箱子，或者用传送带一次性运送。效率瞬间提升百倍。

场景一：寻找两个数组的“共同点”（交集）

普通方法的局限

NumPy 向量化方案：np.in1d 与 布尔索引

1. 检查 arr1 中的每一个元素是否存在于 arr2 中。
2. 返回一个布尔数组（True/False）。
3. 利用这个布尔数组作为“过滤器”（掩膜），从原数组中提取出 True 对应的元素。

PYTHON
import numpy as np

# 创建两个数组
arr1 = np.array([1, 2, 3, 4, 5, 6])
arr2 = np.array([4, 5, 6, 7, 8, 9])

# 1. 制作“寻找器”：判断 arr1 的元素是否在 arr2 中
mask = np.isin(arr1, arr2)

# 此时 mask 是：[False, False, False, True, True, True]

# 2. 使用布尔索引提取“缘分”
common_elements = arr1[mask]

print(f"共同元素: {common_elements}")
# 输出: [4 5 6]

• np.in1d 专门用于一维数组。
• np.isin 是 np.in1d 的多维版本，功能更全面，建议优先使用。

场景二：寻找“独有”元素（差集）

逻辑推导

PYTHON
import numpy as np

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

# 1. 判断 arr1 的元素是否在 arr2 中
mask = np.isin(arr1, arr2)

# 2. 取反，获取“不在”的掩膜
unique_mask = ~mask

# 3. 提取独有元素
unique_elements = arr1[unique_mask]

print(f"arr1 独有的元素: {unique_elements}")
# 输出: [1 2 3]

场景三：查找元素的“位置索引”（对应关系）

核心武器：np.searchsorted

PYTHON
import numpy as np

# 源数组（需要先排序）
arr_source = np.array([10, 20, 30, 40, 50])

# 目标数组（我们要查找这里面的元素在源数组中的位置）
arr_target = np.array([20, 40, 60, 10])

# 1. 对源数组进行排序（如果已经是有序的可跳过）
# 注意：searchsorted 只能在有序数组上工作
sorted_indices = np.argsort(arr_source)  # 获取排序后的索引
sorted_source = arr_source[sorted_indices]

# 2. 使用 searchsorted 查找插入位置
# side='left' 表示查找元素插入的最左侧位置
positions = np.searchsorted(sorted_source, arr_target)

# 3. 处理“未找到”的情况
# searchsorted 对于不存在的元素会返回 len(source)（即插入到最后）
# 我们需要将这些越界的位置标记为 -1
valid_mask = positions < len(sorted_source)

# 进一步检查找到的值是否真的匹配（因为 searchsorted 只保证插入位置）
matches = valid_mask & (sorted_source[positions] == arr_target)

# 4. 还原到原始索引
# 注意：positions 是在排序后数组中的索引，我们需要通过 sorted_indices 映射回原始索引
final_indices = np.where(matches, sorted_indices[positions], -1)

print(f"目标元素在源数组中的原始索引: {final_indices}")
# 输出: [1 3 -1 0]
# 解释: 20在索引1, 40在索引3, 60没找到(-1), 10在索引0

更简单的替代方案：广播（Broadcasting）

PYTHON
import numpy as np

arr1 = np.array([1, 2, 5])
arr2 = np.array([10, 1, 20, 2, 30])

# 广播：arr1 变成列向量，arr2 变成行向量，生成一个比较矩阵
# 结果是一个 (len(arr1), len(arr2)) 的二维数组
comparison_matrix = arr1[:, np.newaxis] == arr2

# 沿着 axis=1（行方向）查找第一个 True 的位置
# axis=1 表示在每一行中查找
indices = np.argmax(comparison_matrix, axis=1)

# 修正未找到的情况
# argmax 在全 False 时会返回 0，这会导致错误，所以需要额外处理
found = np.any(comparison_matrix, axis=1)
final_indices = np.where(found, indices, -1)

print(f"对应索引: {final_indices}")
# 输出: [1 3 -1]

• searchsorted：适合大数组，速度快，但需要排序。
• 广播（Broadcasting）：代码直观，不需要排序，但会消耗较多内存生成矩阵（时间复杂度 O(N*M)）。

性能对比：向量化的威力

总结与最佳实践

1. 寻找交集：使用 np.isin(arr1, arr2) 配合布尔索引。
2. 寻找差集：使用 ~np.isin(arr1, arr2)。
3. 寻找索引：如果数据量大，先排序再用 np.searchsorted；如果数据量小且内存充足，广播机制（== 比较）更直观。
4. 避免循环：时刻提醒自己，Python 的 for 循环在处理密集型数据时是性能杀手，NumPy 的向量化操作是你的最佳伙伴。