Fast way to find index of array in array of arrays

arrays, multidimensional-array, numpy, python, search

Solution

This is Jaime's idea, I just love it:

import numpy as np

def asvoid(arr):
    """View the array as dtype np.void (bytes)
    This collapses ND-arrays to 1D-arrays, so you can perform 1D operations on them.
    https://stackoverflow.com/a/16216866/190597 (Jaime)"""    
    arr = np.ascontiguousarray(arr)
    return arr.view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1])))

def find_index(arr, x):
    arr_as1d = asvoid(arr)
    x = asvoid(x)
    return np.nonzero(arr_as1d == x)[0]


arr = np.array([[  1,  15,   0,   0],
                [ 30,  10,   0,   0],
                [ 30,  20,   0,   0],
                [1, 2, 3, 4],
                [104, 139, 146,  75],
                [  9,  11, 146,  74],
                [  9, 138, 146,  75]], dtype='uint8')

arr = np.tile(arr,(1221488,1))
x = np.array([1,2,3,4], dtype='uint8')

print(find_index(arr, x))

yields

[      3      10      17 ..., 8550398 8550405 8550412]

The idea is to view each row of the array as a string. For example,

In [15]: x
Out[15]: 
array([^A^B^C^D], 
      dtype='|V4')

The strings look like garbage, but they are really just the underlying data in each row viewed as bytes. You can then compare `arr_as1d == x` to find which rows equal `x`.

There is another way to do it:

def find_index2(arr, x):
    return np.where((arr == x).all(axis=1))[0]

but it turns out to be not as fast:

In [34]: %timeit find_index(arr, x)
1 loops, best of 3: 209 ms per loop

In [35]: %timeit find_index2(arr, x)
1 loops, best of 3: 370 ms per loop

Problem

Suppose I have a numpy array of arrays of length 4: ``` In [41]: arr Out[41]: array([[ 1, 15, 0, 0], [ 30, 10, 0, 0], [ 30, 20, 0, 0], ..., [104, 139, 146, 75], [ 9, 11, 146, 74], [ 9, 138, 146, 75]], dtype=uint8) ``` I want to know: - Is it true that `arr` includes `[1, 2, 3, 4]`? - If it true what index of `[1, 2, 3, 4]` in `arr`? I want to find out it as fast as it possible. Suppose `arr` contains 8550420 elements. I've checked several methods with `timeit`: - Just for checking without getting index: `any(all([1, 2, 3, 4] == elt) for elt in arr)`. It tooks 15.5 sec in average on 10 runs on my machine `for`-based solution: `for i,e in enumerate(arr): if list(e) == [1, 2, 3, 4]: break` It tooks about 5.7 secs in average Does exists some faster solutions, for example numpy based?

Original source

Related problems