Set Trailing Zeros to X FAST - performance

Question

I have a 2D array of integers. There may be 0s in the middle of each row, and each row ends with some number of trailing 0s.
How do I set all the trailing zeros to some integer X?
import numpy as np
def generateTestData(N, K, INDEX_SIZE):
# Start with a flat array to easily place zeros inside
data1 = np.random.randint(0, INDEX_SIZE, N*K)
# Add zeros at random locations
idx = np.random.randint(0, N*K, int(N*K/3))
data1[idx] = 0
# Make data1 a (N,K) array
data1 = np.reshape(data1, (N,K))
# Add trailing zeros
for i in range(N):
data1[i,np.random.randint(0,K):] = 0
return data1
if __name__=='__main__':
N = 10000; K = 150; INDEX_SIZE = 500; X = -1
# Test data
data1 = generateTestData(N, K, INDEX_SIZE)
# Save a copy for the test
data2 = np.copy(data1)
for i in range(N):
for j in reversed(range(K)):
if data1[i,j] == 0:
data1[i,j] = X
else:
break
# Faster code here on 'data2'
# ...
def diff(a,b):
return np.mean(np.abs(a-b))
# Verification:
print('Diff(data1,data2) = '+str(diff(data2,data1)))

Answer 1

Here's a vectorized solution leveraging broadcasting -
def replace_trailing_num(a, compare_val=0, assign_val=-1):
idx = a.shape[1] - (a[:,::-1]!=compare_val).argmax(axis=1)
idx[(a==compare_val).all(1)] = 0
mask = np.arange(a.shape[1]) >= idx[:,None]
a[mask] = assign_val
return a
Sample run -
In [60]: a
Out[60]:
array([[2, 3, 0, 4, 6, 0, 0],
[0, 5, 8, 0, 0, 0, 0]])
In [61]: replace_trailing_num(a, compare_val=0, assign_val=-1)
Out[61]:
array([[ 2, 3, 0, 4, 6, -1, -1],
[ 0, 5, 8, -1, -1, -1, -1]])
Alternatively, we can use np.minimum.accumulate to get the mask -
mask = np.minimum.accumulate(a[:,::-1]==compare_val,axis=1)[:,::-1]
If you aren't looping enough or the number of columns is a comparatively bigger number than number of rows, a loopy one based on slicing might be better, one of which is listed below -
def replace_trailing_num_loopy(a, compare_val=0, assign_val=-1):
idx = (a[:,::-1]!=compare_val).argmax(axis=1)
for i,c in enumerate(idx):
a[i,-c:] = assign_val
return a