Indices of intersection of two arrays

Suppose I have two arrays:

a = [2, 4, 6, 7, 10, 11]
b = [6, 7, 10, 11, 13, 15, 17, 19, 21, 23]

How do I get the indices of values in b that are in a and vice versa? I know intersection() returns the values, but I’m after the indices of the two arrays. Right now I’m doing

indices = [findall(x->x==i,a) for i in b]
10-element Array{Array{Int64,1},1}:
 [3]
 [4]
 [5]
 [6]
 [] 
 [] 
 [] 
 [] 
 [] 
 []

which doesn’t seem optimal. Any ideas? Thanks in advance!

You could write a function that loops through the two sorted vectors pretty easily.

1 Like
findall(x->x in b, a)

or

findall(in(b),a)
12 Likes

Perfect, thanks!

I am sharing a routine that efficiently loops through the two sorted vectors, in case someone stumbles upon this thread in the future.

The complexity is O(n+m), where n = length(x) and m = length(y).

function sorted_intersect(x::AbstractVector, y::AbstractVector)

  @assert issorted( x )
  @assert issorted( y )

  idx_common_x = Vector{Int}()
  idx_common_y = Vector{Int}()
  i,j = firstindex(x), firstindex(y)
  @inbounds while i <= lastindex(x) && j <= lastindex(y)
    if x[i] == y[j]
      push!( idx_common_x, i )
      push!( idx_common_y, j )
      i += 1
      j += 1
    elseif x[i] > y[j]
      j += 1
    else
      i += 1
    end
  end

  idx_common_x, idx_common_y

end
4 Likes

:grinning: It would be interesting to benchmark against the other suggestions, for reference.

4 Likes

You may get another 2X speedup by avoiding push! inside the loop and using BitArrays to store common indices flag instead. Although push! is quite efficient in Julia, using it inside hot loops can be expensive. The speedup from this change can vary depending on the ratio of common indices.

function sorted_intersect2(x::AbstractVector, y::AbstractVector)
    @assert issorted( x ) 
    @assert issorted( y ) 

    idx_common_x = falses(length(x))
    idx_common_y = falses(length(y))
    i, j = firstindex(x), firstindex(y)
    while i <= lastindex(x) && j <= lastindex(y)
        if x[i] < y[j]
            i += 1
        elseif y[j] < x[i]
            j += 1
        else
            idx_common_x[i] = true
            idx_common_y[j] = true
            i += 1
            j += 1
        end
    end
    (1:length(x))[idx_common_x], (1:length(y))[idx_common_y]
end

function sorted_intersect(x::AbstractVector, y::AbstractVector)
    @assert issorted( x )
    @assert issorted( y )

    idx_common_x = Vector{Int}()
    idx_common_y = Vector{Int}()
    i,j = firstindex(x), firstindex(y)
    while i <= lastindex(x) && j <= lastindex(y)
        if x[i] == y[j]
            push!( idx_common_x, i )
            push!( idx_common_y, j )
            i += 1
            j += 1
        elseif x[i] > y[j]
            j += 1
        else
            i += 1
        end
    end
    idx_common_x, idx_common_y
end

a = [1:1000;]    #[2, 4, 6, 7, 10, 11]
b = [600:1850;]  #[6, 7, 10, 11, 13, 15, 17, 19, 21, 23]
@assert sorted_intersect(a, b) == sorted_intersect2(a, b)
@btime (findall(in($b),$a) , findall(in($a),$b))
  8.767 μs (10 allocations: 15.12 KiB)
@btime sorted_intersect($a, $b)
  6.700 μs (10 allocations: 15.12 KiB)
@btime sorted_intersect2($a, $b)
  3.263 μs (6 allocations: 7.12 KiB)
4 Likes

Thank you for benchmarking the different solutions.
I am re-running your experiments for larger arrays (1 million elements)

julia> a = sort( rand(1:1_000_000_000, 1_000_000) );
julia> b = sort( rand(1:1_000_000_000, 1_000_000) );
julia> @btime (findall(in($b),$a) , findall(in($a),$b))
  21.785 ms (20 allocations: 32.78 KiB)
julia> @btime sorted_intersect($a, $b)
  12.730 ms (20 allocations: 32.78 KiB)
julia> @btime sorted_intersect2($a, $b)
  11.034 ms (8 allocations: 260.47 KiB)

Version 2 is indeed slightly faster, with a small overhead in memory.

  • Version 1 will allocate memory equal to the size of the intersection, while
  • Version 2 will allocate m + n bytes of memory.

Depending on the application, the user may decide which is one is more appropriate.

1 Like

Indeed, as the vector size increases, memory allocations happen less and less often and the benefit from using a BitArray becomes less evident after some threshold. Many algorithms follow this pattern, but the problem is that it is difficult to predict such a threshold, especially on different hardware. I remember applying these tricks to some algorithms like in Primes.jl, and around 10^9 vector lengths, push! becomes more efficient. I imagine push! will become more smarter in the future or hardware will be able to handle this and we won’t need such tricks.

2 Likes