Numpy 2d array, obtain indexes of rows where specified column indexes equal 1

user_5124717 0
Viewed 832

I have a 2d numpy array like so, which will only ever have 0, 1 values.

a = np.array([[1, 0, 1, 0],  # Indexes 0 and 2 == 1
             [0, 1, 1, 0],   # Indexes 1 and 2 == 1
             [0, 1, 0, 1],   # Indexes 1 and 3 == 1
             [0, 1, 1, 1]])  # Indexes 1, 2, and 3 == 1

What I'd like to do is to get the indices of every row where a pair of column indexes passed are all equal to 1.

For example, if the function doing this is get_rows, get_rows(a, [1, 3]), should return [2, 3], because the rows at indexes 2 and 3 have column indexes 1 and 3 equal to 1. Similarly, get_rows(a, [1, 2]) should return [1, 3].

I know how to do this in a Pandas dataframe, but I'd like to stick to using pure numpy for this one. I tried using np.where in some form like 

np.where( ((a[i1 - 1] == 1) & (a[i2 - 1] == 1) ))

but this doesn't seem to give me what I want, and wouldn't work for varying number of passed indices.

numpy arrays python
2 Answers

I think you are looking for this:

col_idx = [1, 2]
np.where(a[:,col_idx].all(axis=1))[0]

You can use ANY column indexes you want to pass to it. It is quite self explanatory that extracts columns and searches for rows with all 1s in there using np.where.

EDIT: per @Mad Physicist recommendation, here is another similar solution:

np.flatnonzero(a[:,col_idx].all(axis=1))

Output example for your input:

[1 3]
user_49759810

Solution

Try this.

Code - Logic

  • Find row and column indices where the given columns (target_col_index) have 1.
  • Select only those rows where column indices match the target_col_index.
  • Narrow down the result by checking the which row indices have results matching to the same number of columns as was in target_col_index.
import numpy as np

target_col_index = [1,2]
target_row_index = get_row_index(a, target_col_index)
print(target_row_index)

## Output
# [1,3]

## Other cases tested
# test_col_indexes = [ [0,1], [0,2], [0,3], [1,2], [1,3], [2,3], [0,1,3], [1,2,3] ]
# returned_row_indexes = [ [], [0], [], [1,3], [2,3], [3], [], [3] ]

Code - Custom Function

def get_row_index(arr, target_col_index=None):
    if target_col_index is None:
        return None
    else:
        row_index, col_index = np.where(arr==1)
        result = row_index[np.isin(col_index, target_col_index)]
        rows, counts = np.unique(result, return_counts=True)
        target_row_index = rows[counts==len(target_col_index)]
        return target_row_index

Dummy Data

a = np.array([[1, 0, 1, 0],  # Indexes 0 and 2 == 1
            [0, 1, 1, 0],   # Indexes 1 and 2 == 1
            [0, 1, 0, 1],   # Indexes 1 and 3 == 1
            [0, 1, 1, 1]])  # Indexes 1, 2, and 3 == 1
user_84748940
Related