CODEWITHSUNDEEP
pythonarraysnumpyindexing
Eric Snyder
Eric Snyder

Reputation: 69

Getting indices of array elements matching a value

I've spent a couple of days on the indexing documentation but haven't found what I am looking for.

Consider this:

import numpy
fac=numpy.asarray(['a','a','a','b','b','c','c','c'])
f_ind = [x for x in range(len(fac)) if fac[x] == 'c']

it returns [5,6,7] like I want. However, it seems like NumPy arrays should have a mechanism for achieving the same thing in a more concise (and efficient?) way. Boolean arrays might be part of the solution:

ba = (fac == 'c')
f_vals = fac[ba]

But that only regurgitates the elements of fac that equal 'c' -- not very helpful.

Any suggestions on how to make this happen using NumPy? Or should I just be happy with what I've got?

Upvotes: 4

Views: 4752

Answers (2)

MSeifert
MSeifert

Reputation: 152785

There are several ways to adress this with NumPy, depending on your needs you could use:

>>> import numpy as np
>>> fac = np.asarray(['a','a','a','b','b','c','c','c'])

  • where:

    >>> np.where(fac == 'c')
(array([5, 6, 7], dtype=int64),)

  • argwhere:

    >>> np.argwhere(fac == 'c')
array([[5],
       [6],
       [7]], dtype=int64)

  • flatnonzero:

    >>> np.flatnonzero(fac == 'c')
array([5, 6, 7], dtype=int64)

Upvotes: 3

Mrl0330
Mrl0330

Reputation: 140

The function you are looking for is numpy's where.

https://docs.scipy.org/doc/numpy-1.12.0/reference/generated/numpy.where.html

np.where(fac=='c')

Would return a tuple containing an array of indexes that match that value, and the data type.

Upvotes: 0

Related Questions