CODEWITHSUNDEEP
pythonmatlabnumpy
Con Mai
Con Mai

Reputation: 33

Replicating the indices result of Matlab's ISMEMBER function in NumPy?

I have been racking my brain for a solution that is along the lines of this older question. I have been trying to find a Python code pattern that replicates the indices result. For example:

A = [3;4;4;3;6]
B = [2;5;2;6;3;2;2;5]
[tf ix] = ismember(A,B)
>> A(tf)

ans =

     3
     3
     6
>> B(ix(tf))

ans =

     3
     3
     6

What this allows me to do is if there is an array C ordered the same way as B I can now appropriately insert the values of C into a new array D that is ordered the same way as A. I do this mapping of data a lot! I would love for this to work for various data types like strings and datetimes in particular. It seems that numpy's in1d get's me half way there. I'm also open to other Pythonic ideas as well! 

D(tf) = C(ix(tf))

Thank you!

Upvotes: 3

Views: 1674

Answers (1)

unutbu
unutbu

Reputation: 879251

import numpy as np

A = np.array([3,4,4,3,6])
B = np.array([2,5,2,6,3,6,2,2,5])

def ismember(a, b):
    # tf = np.in1d(a,b) # for newer versions of numpy
    tf = np.array([i in b for i in a])
    u = np.unique(a[tf])
    index = np.array([(np.where(b == i))[0][-1] if t else 0 for i,t in zip(a,tf)])
    return tf, index

tf,ix=ismember(A,B)
print(tf)
# [ True False False  True  True]
print(ix)
# [4 0 0 4 5]
print(A[tf])
# [3 3 6]
print(B[ix[tf]])
# [3 3 6]

Upvotes: 5

Related Questions