Cycle through NP array with different lengths of each row

Question

I have one np array with 10 rows and 5 columns. I also have an array with 20 rows and a single column.

a = [[15,2,1,0,0],
     [8,12,0,0,0],
     [4,12,10,9,0],
     [3,0,0,0,0],
     [19,7,0,0,0],
     [13,15,4,0,0],
     [0,0,0,0,0],
     [11,4,0,0,0],
     [7,0,0,0,0,0],
     [10,6,8,4,0]]



b =  [8,1,6,4,9,3,5,6,11,14,5,4,33,7,9,15,7,3,19,3]

For array a, once a value in a row is zero, any value after that will be zero.

I would like to cycle through each row in a and find the nth value in b, and then store that in a 10x5 array called c. Any zero in a should also be a zero in c.

For example, the first two rows in c would be:

[[9,1,8,0,0],
 [6,4,0,0,0]]

Thank you!

mozway · Accepted Answer · 2023-04-05T14:17:25.837

You can use indexing:

# output with the shape of a
c = np.zeros_like(a)

# non-zero positions
m = a != 0

# assign values
c[m] = b[a[m]-1]

Shorter alternative with numpy.where (but maybe less efficient if you have a lot of zeros):

c = np.where(a!=0, b[a-1], 0)

Output:

array([[ 9,  1,  8,  0,  0],
       [ 6,  4,  0,  0,  0],
       [ 4,  4, 14, 11,  0],
       [ 6,  0,  0,  0,  0],
       [19,  5,  0,  0,  0],
       [33,  9,  4,  0,  0],
       [ 0,  0,  0,  0,  0],
       [ 5,  4,  0,  0,  0],
       [ 5,  0,  0,  0,  0],
       [14,  3,  6,  4,  0]])

Reproducible input:

a = np.array([[15,  2,  1,  0,  0],
              [ 8, 12,  0,  0,  0],
              [ 4, 12, 10,  9,  0],
              [ 3,  0,  0,  0,  0],
              [19,  7,  0,  0,  0],
              [13, 15,  4,  0,  0],
              [ 0,  0,  0,  0,  0],
              [11,  4,  0,  0,  0],
              [ 7,  0,  0,  0,  0],
              [10,  6,  8,  4,  0]])


b =  np.array([ 8,  1,  6,  4,  9,  3,  5,  6, 11, 14,  5,  4, 33,  7,  9, 15,  7,
                3, 19,  3])

timings

On (10000, 5) with mostly non-zeros

# indexing
203 µs ± 6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

# numpy.where
97.2 µs ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

On (10000, 10005) with only the first 5 columns as non-zero

# indexing
2.2 s ± 384 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# numpy.where
4.73 s ± 680 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Cycle through NP array with different lengths of each row

1 Answers1

timings