Get multiple results as output for one condition on numpy array

Viewed 85

I am looking for a way to get multiple results for one condition of a numpy array. Currently, I run the same where condition twice on the same numpy array, to get two new arrays with different output.

import numpy as np
np_array = np.random.rand(3,3)
entry_array = np.where(np_array>=0.5,1,0)
mod_array = np.where(np_array>=0.5,np_array,0)

Does a possibility exists where I can fill the two new arrays entries, mod_array with only one loop over np_array. Additionally, I am also looking for a solution for non-array output, like the following.

entries = np.where(np_array>=0,1,0).sum()
sum_entries = np.where(np_array>=0,np_array,0).sum()

I am especially interested in an efficient way, as my arrays have more than 100 million entries.

1 Answers

Few similar ideas:

def m1(np_array):
  entry_array = np.where(np_array>=0.5,1,0)
  mod_array = np.where(np_array>=0.5,np_array,0)
  return mod_array, entry_array

def m2(np_array):
  mod_array = np.where(np_array>=0.5,np_array,0)
  entry_array = np.array(mod_array, dtype=bool)+0
  return mod_array, entry_array

def m3(np_array):
  entry_array = np_array<0.5
  np_array[entry_array] = 0
  return np_array, 1-entry_array

Comparison: m1 seems slightly slower (up to 2x) than m2 and m3.

in_ = [np.random.rand(n,n) for n in [10,100,1000,10000,20000]]

enter image description here

Related