trim np arrays according to a list of starting points-CodePudding

I have a table, represented by an np.array like the following:

A = [[12,412,42,54],
     [144,2,42,4],
     [2,43,22,10]]

And a list that contains the desired starting point of each row in A:

L=[0,2,1]

The desired output would be:

B = [[12,412,42,54],
     [42,4,np.nan,np.nan],
     [43,22,10,np.nan]]

Edit
I prefer to avoid using a for-loop for obvious reasons.

CodePudding user response：

Try compare the L with column index, then use boolean set/get items:

# convert A to numpy array for advanced indexing
A = np.array(A)
ll = A.shape[1]


keep = np.arange(ll) >= np.array(L)[:,None]

out = np.full(A.shape, np.nan)
out[keep[:,::-1]] = A[keep]
print(out)

Output:

[[ 12. 412.  42.  54.]
 [ 42.   4.  nan  nan]
 [ 43.  22.  10.  nan]]

CodePudding user response：

My guess would be that a vectorized approach for this would be less efficient than explicit looping, because the result is fundamentally a jagged array, which NumPy does not support well.

However, a loop-based solution is simple, that can be made faster with Numba's