从Python转Julia：如何更高效获取数组b中存在于a的元素？

你从Python转Julia，想要提取数组b中所有存在于数组a的元素，现有写法虽然能实现功能，但在性能和简洁性上还有优化空间，以下是几种更好的实现方式：

现有代码的问题 #

你当前的Julia代码：

indices_b_in_a = findall(ele -> ele in a, b);
b_in_a = b[indices_b_in_a];

每次判断ele in a都会遍历整个a数组，当a元素较多时，时间复杂度是O(n*m)（n是b长度，m是a长度），性能会明显下降。

优化方案 #

1. 用集合（Set）优化查询（最通用高效） #

把a转换成Set，集合的成员查询是O(1)时间复杂度，整体复杂度降到O(n+m)，代码也更简洁：

a = [1, 2, 3, 4];
b = [7, 8, 2, 3, 5];
a_set = Set(a);
# 列表推导式写法
b_in_a = [ele for ele in b if ele in a_set];
# 或者布尔索引写法，更贴近numpy的风格
b_in_a = b[in.(b, Ref(a_set))];

两种写法都能得到[2, 3]，且保留b中元素的原始顺序和重复项（比如如果b是[3,2,2,5]，结果会是[3,2,2]）。

2. 用intersect快速求交集（适合不关心顺序/去重的场景） #

如果不需要保留b中的原始顺序，也不需要保留重复元素，可以直接用Julia内置的intersect函数：

b_in_a = intersect(b, a);

注意：intersect会返回去重且按元素大小排序的结果，比如b = [3,2,2,5]时，结果会是[2,3]，这点和numpy的isin行为不同，需要根据需求选择。

3. 排序+二分查找（适合大型数值数组） #

如果处理的是超大数值数组，先对a排序，再用二分查找判断元素是否存在，时间复杂度是O(n log m)，性能也很出色：

a_sorted = sort(a);

# 封装一个判断元素是否在有序数组里的函数
function in_sorted(ele, sorted_arr)
    idx = searchsortedfirst(sorted_arr, ele)
    return idx <= length(sorted_arr) && sorted_arr[idx] == ele
end

b_in_a = [ele for ele in b if in_sorted(ele, a_sorted)];

这种方法在数值类型数组上的性能和Set方法接近，但对于不可排序的类型（比如自定义结构体）不适用。

性能总结 #

• 小数组场景：几种方法差异不大，选你觉得最易读的就行；
• 中大型数组：Set方法是通用首选，排序二分法适合数值数组；
• 去重+排序需求：直接用intersect最简洁。

内容的提问来源于stack exchange，提问作者Jared