Skip to content
Advertisement

Working with 2 arrays to populate a third one in numpy

I am trying to work with two arrays in a certain way in python. Lets say

A = np.array([5, 10, 30, 50])
B = np.array([2, 3, 4, 5])

now I have a target array C = [2, 7, 15, 25, 40]. I want to find a output value (say Y) for each element in C (say x) with conditions:

  • If x < A[0] then Y = B[0]
  • If A[i] < x < 2*A[i] < A[i+1] or A[i] < x < A[i+1] < 2*A[i] then Y=x
  • If A[i] < 2*A[i] < x < A[i+1] then Y = B[i+1]
  • If x > A[end] then Y = x

i is the maximum possible index satisfying the conditions

So, output becomes Y = [2, 7, 15, 4, 40]

This is what I have currently

A = np.array([5, 10, 30, 50])
B = np.array([2, 3, 4, 5])
C = np.array([2, 7, 15, 25, 40])

diffA = np.diff(A)

diffIdx = np.squeeze(np.where(diffA > A[0:len(A)-1]+1e-6))

Y = np.array([])

for x in C:
  idx = np.searchsorted(A, x)
  if idx == 0:
    Y = np.append(Y, B[0])
  elif np.any(diffIdx==idx-1) & (x>2*A[idx-1]+1e-6):
    Y = np.append(Y, B[idx])
  else:
    Y = np.append(Y, x)

This seems to work, and IPython console %timeit shows it takes 24.7 ms, but I was wondering if there is a better or faster method of doing the same.

Advertisement

Answer

Check this question Most efficient way to map function over numpy array . You can create an array of indices np.arange(C.size), make a function with all the logic of combining A, B and C element getting them by index and depending on the size of your arrays choose the fastest method to apply the function.

User contributions licensed under: CC BY-SA
1 People found this is helpful
Advertisement