CODEWITHSUNDEEP
pythonarraysnumpy
Forklift17
Forklift17

Reputation: 2477

Replacing numpy 2D array elements below an index array

I searched prior questions and didn't find one that quite fits this.

I am setting the value of all cells in A's jth column to 0 whose row indeces are less than m[j]:

import numpy as np
n = 8
A = np.ones([n,n])
m = np.array([1,1,1,2,2,2,3,4])
for j in range(n):
    for i in range(m[j]):
        A[i,j] = 0

How do I rewrite that script without 'for' loops?

Upvotes: 3

Views: 76

Answers (2)

cs95
cs95

Reputation: 402263

You can use broadcasted NumPy comparison—

>>> A[np.arange(n)[:, None] < m] = 0
>>> 
array([[0., 0., 0., 0., 0., 0., 0., 0.],
       [1., 1., 1., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 1., 0., 0.],
       [1., 1., 1., 1., 1., 1., 1., 0.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.]])

Here, calling [:, None] augments the shape of np.arange(n) such that the comparison < is broadcasted across every element of m, for each item in the range. This generates a boolean mask of the same shape as A which is then used to set values to 0.

Note—If A is guaranteed to be an array of ones, I would recommend Divakar's solution, which is very similar to this.

Upvotes: 4

Divakar
Divakar

Reputation: 221514

Use outer-broadcasted comparison against the ranged array of length n to get directly get A -

A = (m <=  np.arange(n)[:,None]).astype(float)

To boost up the performance, we could use the type conversion to lower precision ones, like .astype(np.int) or even .astype(np.uint8). If you are okay with a boolean array, skip the type conversion altogether.

For int dtypes, we can also simply view into it, that saves us some memory -

(m <=  np.arange(n)[:,None]).view(np.int8)

Upvotes: 3

Related Questions