How can I find matching elements and indices from two arrays?

Viewed 2538

For example,

a = [1, 1, 2, 4, 4, 4, 5, 6, 7, 100]
b = [1, 2, 2, 2, 2, 4, 5, 7, 8, 100]

I can find the matching elements using:

np.intersect1d(a,b)

Output:

array([  1,   2,   4,   5,   7, 100])

Then, how can I get the indices of matched elements in arrays a and b, respectively ?

There is a function in IDL as "match" - https://www.l3harrisgeospatial.com/docs/match.html

Is there a similar function in Python?

5 Answers

Use return_indices in numpy.intersect1d:

intersect, ind_a, ind_b = np.intersect1d(a,b, return_indices=True)

Output:

intersect
# array([  1,   2,   4,   5,   7, 100])
ind_a
# array([0, 2, 3, 6, 8, 9], dtype=int64)
ind_b
# array([0, 1, 5, 6, 7, 9], dtype=int64)

Which can then be reused like:

np.array(a)[ind_a]
np.array(b)[ind_b]

array([  1,   2,   4,   5,   7, 100])

you can track indexes using enumerate as follows:

a = [1, 1, 2, 4, 4, 4, 5, 6, 7, 100]
b = [1, 2, 2, 2, 2, 4, 5, 7, 8, 100]
#    0  1  2  3  4  5  6  7  8   9

print([i for i,x in enumerate(zip(a,b)) if x[0] == x[1]])
[0, 2, 5, 6, 9]

so what's going on here?!

We're taking advantage of the amazing enumerate function! This function generates a tuple for each element in an iterable, with the first element being the enumeration (or the index in this case) and the second being the iterable.

Heres what the enumeration of zip(a,b) looks like

[(0, (1, 1)), (1, (1, 2)), (2, (2, 2)), (3, (4, 2)), (4, (4, 2)), (5, (4, 4)), (6, (5, 5)), (7, (6, 7)), (8, (7, 8)), (9, (100, 100))]

# lets look a little closer at one element
(0, (1, 1))
# ^     ^
# index iterable

From there on it's simple! Unpack the iterable and check if both elements are equal, and if they are then use append the enumeration # to a list!

Use range like so:

matching_idxs = [idx for idx in range(len(a)) if a[idx] == b[idx]] 
print(matching_idxs)
# [0, 2, 5, 6, 9]

Using enumerate and zip:

a = [1, 1, 2, 4, 4, 4, 5, 6, 7, 100]
b = [1, 2, 2, 2, 2, 4, 5, 7, 8, 100]

output = [(i, x) for i, (x, y) in enumerate(zip(a, b)) if x == y]

print(output)

[(0, 1), (2, 2), (5, 4), (6, 5), (9, 100)]

This results in a list of tuples being (index, value)

this is very straightforward with the help of zip function to loop two lists in parallel:

>>> count = 0
>>> indices = []
>>> for x, y in zip(a, b):
    if x == y:
        indices.append(count)
    count += 1

    
>>> indices
[0, 2, 5, 6, 9]
Related