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 love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With