Searching numpy array for for pattern

Question

I'd like to find a value in a numpy array given a search pattern. For instance for the given array a, I want to retrieve a result of 1 when using the search pattern s because 1 is the element at index 0 of a[:,1] (=array([1, 0, 0, 1])) and the elements of a[1:,1] match s (i.e. (a[1:,1] == s).all() == True => return a[0,1]).

Another example would be s=[1, 0, 1] for which I would expect a search result of 2 (match at 4th column starting (1-based)). 2 would also be the search result for s=[2, 0, 0], etc.

>>> import numpy as np
>>> a = np.asarray([[0, 1, 2, 2, 2, 2, 2, 2], [0, 0, 1, 1, 2, 2, 3, 3], [0, 0, 0, 0, 0, 0, 0, 0], [0, 1, 0, 1, 0, 1, 0, 1]])
>>> a
array([[0, 1, 2, 2, 2, 2, 2, 2],
       [0, 0, 1, 1, 2, 2, 3, 3],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 1, 0, 1, 0, 1]])
>>> s = np.asarray([0, 0, 1])

I came up with a[0, np.where((a[1:,:].transpose() == s).all(axis=-1))[0][0]], but thought there must be something more elegant...

Additionally, it would be great if I could do this operation with one call on multiple search patters, so that I retrieve the 0-element for which the values of index 1 to index 3 match.


Show source
| numpy   | python   | arrays   2017-01-06 12:01 1 Answers

Answers ( 1 )

  1. 2017-01-06 12:01

    Single search pattern

    Here's one approach with help from broadcasting and slicing -

    a[0,(a[1:] == s[:,None]).all(0)]
    

    Multiple search patterns

    For multiple search patterns (stored as 2D array), we just need to broadcast as before and look for ANY match at the end -

    a[0,((a[1:] == s[...,None]).all(1)).any(0)]
    

    Here's a sample run -

    In [327]: a
    Out[327]: 
    array([[0, 1, 2, 2, 2, 2, 2, 2],
           [0, 0, 1, 1, 2, 2, 3, 3],
           [0, 0, 0, 0, 0, 0, 0, 0],
           [0, 1, 0, 1, 0, 1, 0, 1]])
    
    In [328]: s
    Out[328]: 
    array([[1, 0, 1],
           [2, 0, 0]])
    
    In [329]: a[0,((a[1:] == s[...,None]).all(1)).any(0)]
    Out[329]: array([2, 2])
    
◀ Go back