CODEWITHSUNDEEP
pythonnumpy
Sengiley
Sengiley

Reputation: 289

Find index of a row in numpy array

Given m x n numpy array

X = np.array([
  [1, 2],
  [10, 20],
  [100, 200]
])

how to find index of a row, i.e. [10, 20] -> 1?

n could any - 2, 3, ..., so I can have n x 3 arrays

Y = np.array([
  [1, 2, 3],
  [10, 20, 30],
  [100, 200, 300]
])

so I need to pass a vector of size n, in this case n=3, i.e a vector [10, 20, 30] to get its row index 1? Again, n could be of any value, like 100 or 1000.

Numpy arrays could be big, so I don't want to convert them to lists to use .index()

Upvotes: 1

Views: 5153

Answers (3)

NMZ
NMZ

Reputation: 66

Just in case that the query array contains duplicate rows that you are looking for, the function below returns multiple indices in such case.

def find_rows(source, target):
    return np.where((source == target).all(axis=1))[0]

looking = [10, 20, 30]

Y = np.array([[1, 2, 3],
              [10, 20, 30],
              [100, 200, 300],
              [10, 20, 30]])

print(find_rows(source=Y, target=looking)) # [1, 3]

Upvotes: 4

zoldxk
zoldxk

Reputation: 1

You can make a function as follow:

def get_index(seq, *arrays):
    for array in arrays:
        try:
            return np.where(array==seq)[0][0]
        except IndexError:
            pass

then:

>>>get_index([10,20,30],Y)
1

Or with just indexing:

>>>np.where((Y==[10,20,30]).all(axis=1))[0]
1

Upvotes: 0

akuiper
akuiper

Reputation: 215117

You can use numpy.equal, which will broadcast and compare row vector against each row of the original array, and if all elements of a row are equal to the target, the row is identical to the target:

import numpy as np
np.flatnonzero(np.equal(X, [10, 20]).all(1))
# [1]

np.flatnonzero(np.equal(Y, [10, 20, 30]).all(1))
# [1]

Upvotes: 2

Related Questions