How to find two specific values in a large 6-column array in Python?

user_14951957 0
Viewed 60

I have a 1213x5 column array of rows with between 3 and 5 values (see example below).

Example of array:

2 3 6 nan nan
8 6 2 nan nan
9 8 6 5 nan
9 5 2 1 nan
2 3 4 1 6
6 8 5 3 2

I have two values and want to find the index of the rows where the two values appear. I have created the following code however it takes ~0.2 seconds to complete. This is not ideal since I have to iterate over this thousands of times. I'm new to Python and therefore am still getting used to pythonic code and what the quickest way to do things are.

Current code:

def rowIdx(array, m, n):
    idxList = []
    m = int(m)
    n = int(n)
    for x in range(len(array)):
        if (array[x,0]== m or array[x,0] == n) and (array[x,1] == m or array[x,1] == n):
            idxList.append(x)
        if (array[x,0] == m or array[x,0] == n) and (array[x,2] == m or array[x,2] == n):
            idxList.append(x)
        if (array[x,1] == m or array[x,1] == n) and (array[x,2] == m or array[x,2] == n):
            idxList.append(x)
        if (array[x,0]== m or array[x,0] == n) and (array[x,3] == m or array[x,3] == n):
            idxList.append(x)
        if (array[x,1] == m or array[x,1] == n) and (array[x,3] == m or array[x,3] == n):
            idxList.append(x)
        if (array[x,2] == m or array[x,2] == n) and (array[x,3] == m or array[x,3] == n):
            idxList.append(x)
        if (array[x,0] == m or array[x,0] == n) and (array[x,4] == m or array[x,4] == n):
            idxList.append(x)
        if (array[x,1]== m or array[x,1] == n) and (array[x,4] == m or array[x,4] == n):
            idxList.append(x)
        if (array[x,2] == m or array[x,2] == n) and (array[x,4] == m or array[x,4] == n):
            idxList.append(x)
        if (array[x,3] == m or array[x,3] == n) and (array[x,4] == m or array[x,4] == n):
            idxList.append(x)            
    return idxList

Any help would be much appreciated. Thanks!

data-science python performance
4 Answers

None of the answer presented above are very efficient as they rely on Python for loops and I wouldn't consider them too Pythonic either. In this case you are better of using just numpy functions. Under the hood they are running C code which is heavily optimized for the kind of thing you want to do.

Hopefully, the code below works for you. This is going to be orders of magnitude faster than any of the solutions currently proposed.

import numpy as np

# Example array
array = np.array([[1,2,90,91,90],[90,2,3,91,5],[1,2,3,np.nan,np.nan]])

# Adapt to your values.
# Find rows per value where the element is there. Easy to extend to more than two values
msk_value_1 = (array == 90).any(axis=1)
msk_value_2 = (array == 91).any(axis=1)

both_true = msk_value_1 & msk_value_2

# find the indices
np.where(both_true)

Benchmark

Here are some benchmarks to compare against most of the solutions provided:

# Create a random matrix for inputs
array = np.random.randint(0,1000, (3000, 4))

# solution here
def find_two_values(array, m, n):
    msk_value_1 = (array == int(m)).any(axis=1)
    msk_value_2 = (array == int(n)).any(axis=1)

    both_true = msk_value_1 & msk_value_2

    # find the indices
    return np.where(both_true)

# takes 900 microseconds on my computer
_ = find_two_values(array, 90, 91)

# other solution proposed
def rowIdx(array, m, n):
    idxList = []
    m = int(m)
    n = int(n)
    for rx, aRow in enumerate(array):
        if m in aRow and n in aRow:
            idxList.append(rx)
# takes 9.8 ms on my computer
_ = rowIdx(array, 90, 91)

That is 1ms vs 10ms roughly ten times faster.

user_9385300

I would think that iterating over the rows of the array and checking for the desired values with in should be straightforward

def rowIdx(array, m, n):
    idxList = []
    m = int(m)
    n = int(n)
    for rx, aRow in enumerate(array):
        if m in aRow and n in aRow:
            idxList.append(rx)

enumerate is a standard built-in function, well worth getting comfortable with.

If you're comfortable with it, you can do the same as a list comprehension:

def rowIdx(array, m, n):
    m = int(m)
    n = int(n)
    idxList = [rx for rx, aRow in enumerate(array) if m in aRow and n in aRow]
user_44754680

This is very pythonic and efficient:

def row_index(array,m,n):
    for index, row in enumerate(array):
        if m in row and n in row:
            yield index

If you want to get the result as a list of indices, than just call it like this:

list(row_index(array,m,n))
user_146915330

You can enumerate over list (either numpy.array or regular python list). Instead of checking each element and position, check if element belongs to a list:

def newRowIdx(array, m, n):
    idxList = []
    m = int(m)
    n = int(n)
    for idx, row in enumerate(array):
        if (m in row and n in row):
            idxList.append(idx)

This method takes 0.372s, while yours took 0.976s (checking in an online environment, so all are slow) but the improvement is noticible.

user_66456240
Related