I need to do logical iteration over numpy array, which's values depend on elements of other array. I've written code below for clarifying my problem. Any suggestions to solve this problem without for loop?
Code
a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = np.zeros_like(b)
c[0] = 150
for i in range(1, c.size):
if a[i] == "b":
c[i] = c[i-1]
else:
c[i] = b[i]
Here's an approach using a combination of np.maximum.accumulate and np.where to create stepped indices that are to be stopped at certain intervals and then simply indexing into b would give us the desired output.
Thus, an implementation would be -
mask = a!="b"
idx = np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))
out = b[idx]
Sample step-by-step run -
In [656]: # Inputs
...: a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
...: b = np.array([150, 154, 147, 126, 148, 125])
...:
In [657]: mask = a!="b"
In [658]: mask
Out[658]: array([ True, False, True, True, False, True], dtype=bool)
# Crux of the implmentation happens here :
In [696]: np.where(mask,np.arange(mask.size),0)
Out[696]: array([0, 0, 2, 3, 0, 5])
In [697]: np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))
Out[697]: array([0, 0, 2, 3, 3, 5])# Stepped indices "intervaled" at masked places
In [698]: idx = np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))
In [699]: b[idx]
Out[699]: array([150, 150, 147, 126, 126, 125])
You could use a more vectorized approach Like so:
np.where(a == "b", np.roll(c, 1), b)
np.where will take the elements from np.roll(c, 1) if the condition is True or it will take from b if the condition is False. np.roll(c, 1) will "roll" forward all the elements of c by 1 so that each element refers to c[i-1].
These type of operations are what make numpy so invaluable. Looping should be avoided if possible.
If you don't need to wrap around the margin there is a very simply solution:
a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = b.copy() #removes necessity of else case
c[a[:-1]=='b'] = c[a[1:]=='b']
or equally:
a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = b.copy() #removes necessity of else case
mask = a == 'b'
c[mask[:-1]] = c[mask[1:]]
If you want to wrap around the margin (a[0]=='b') then it gets a little more complicated, you either need to use roll or catch this case first with and if.
