Find Values in an array

This script demonstrates how to find values in a numpy array and return the cell locations and the value. Obviously this can be modified to suit one's purposesl

'''FindValsInArray.py

   Demos how to query a ndarray of data for cell locations meeting certain conditions.
   Returns the row, column and value of the array locations (ie cells) meeting the condition
'''

import numpy as np

def find_in_array(an_array, a_query):
  '''find the locations of cells in an ndarray whose values meet some condition'''
  rc_IDs = np.indices(an_array.shape)        # create row, col index arrays (could be 1D, 2D etc)
  located = rc_IDs[:,a_query]                # query the array
  pnts = np.hsplit(located, len(located[0])) # get the r, c locations, separately
  locations = []
  for pnt in pnts:                           # build the output as a list of
    r, c = pnt.flatten().tolist()             # X, Y, val  (ie col, row, val)
    a_val = an_array[r, c]
    locations.append( [r, c, a_val] )
  return locations

if __name__ == "__main__":
  ''' The find_in_array function is used to find conditions in an ndarray
      that correspond to some condition.  The array and condition is passed
      to it ie
     an_array = some array of any shape
     example queries:
       a_query = an_array >= 22                               # a single condition
       a_query = np.logical_or(an_array >= 22, an_array <= 4) # multiple conditions
     see logic functions in numpy, eg.  logical_and, logical_not
  '''
  data = [[1,2,3,4,5],
          [6,7,8,9,10],
          [11,12,13,14,15],
          [16,17,18,19,20],
          [21,22,23,24,25]
         ]
  an_array = np.array(data)                              #the array to query
  #a_query = an_array >= 22                               #sample query 1
  a_query = np.logical_or(an_array >= 22, an_array <= 4) #sample query 2
  
  locations = find_in_array(an_array, a_query)           #do the work
  print "Input data as an example"
  print "["
  for i in data:
    print i
  print "]"
  print "Querying for an_array >=22 and an_array <=4"
  print "Row, column and values are listed below"
  print "\nrow, col, value)"
  for i in locations:
    print i