Python Forum

Consider the following code:

import numpy as np
y_i=np.linspace(0,1,11) #0,0.1,0.2,...,1
y=np.random.uniform(0,1) #random point in [0,1)
i=np.argmax(y_i>=y) #returns the index of the smallest number in y_i greater than or equal to y

If y is an array of random points, and one would like to attain all the indices in y_i for each of these points, how would one go about doing that using np.argmax or any other one-liner if it exists?

Current solution:

d=2
y=np.random.uniform(0,1,d)
i=np.zeros(d)
for l in range(0,d):
  i[l]=np.argmax(y_i>=y[l])

Your solution is quite slow because you are using pure Python loops. For this problem
complete vectorized solution exists:

np.argmax(np.repeat(y_i[:, np.newaxis], len(y), axis=1) >= y, axis=0)

It is expected that y is a numpy array of shape (len(y), ).

Suppose y_i is a matrix such that

y_i=np.linspace(np.zeros(d),np.ones(d),m+1)  

That is, each column is an array given by np.linspace(0,1,m+1) and the entries between 0 and 1, i.e. the first and last row, might be assigned different values in a later statement. Then for y being an array of random points, I'd like to know, for each column of the matrix y_i the index of the number greater than or equal to a random point in y. In other words, I'd like to get a matrix of indices of shape (len(y),d).

Current solution:

y_i=np.linspace(np.zeros(d),np.ones(d),m+1) 
y=np.random.uniform(0,1,d)                      
ind=np.zeros([len(y),d]) 
for k in range(d):
  k_i=y_i[:,k]
  i=np.argmax(np.repeat(k_i[:, np.newaxis], len(y), axis=1) >= y, axis=0)
  ind[:,k]=i

I realize that what I want is not quite what I have said, so I'm restating it:

Suppose

y_i=np.linspace(np.zeros(d),np.ones(d),m+1)

and the values between 0 and 1 (the first and last row) are changed in a later statement of the code. Then, for

y=np.random.uniform(0,1,d)

I'd like to find the index of the number in the i'th column of y_i greater than or equal to the i'th entry of y.

Current solution:

y_i=np.linspace(np.zeros(d),np.ones(d),m+1)
y=np.random.uniform(0,1,d)                      
ind=np.zeros(d) 
for k in range(d):
   k_i=y_i[:,k]
   i=np.argmax(k_i >= y[k])
   ind[k]=i