Return top “N” elements from an array

top_n returns a mask = [True, False, True, False, False ...] with “True” for top n values. The mask is passed into an array as index to get “True” values.

import numpy as np
from scipy.stats import rankdata

def top_n(list_array, n = 1):

  """
  Returns a boolean mask with "True" for greatest "n" number of values
  """
  np_array = np.array(list_array)
  # creating a mask
  mask = np.zeros(len(np_array.flatten()), dtype=bool)
  r =rankdata(np_array, method ="dense")
  # rank matrix with highest value =1
  r=(r.max()+1)-r
  for index, val in enumerate(r):
    if  val <= (n):
	mask[index] = True
  return mask.reshape(np_array.shape)

boolean_filter will return a list where boolean is true.

def boolean_filter(b_list, boolean):
  """
  This function returns values in b_list where the boolean is true
  """
  return [item for i, item in enumerate(b_list) if boolean[i]==True]