Efficiently zero elements of numpy array using a boolean mask

user_1973207 0
Viewed 289

So I have created a super-slow version:

arr = np.arange(3*9).reshape((3, 9))
print(arr)
to_black = np.random.choice(a=[False, True], size=9)
for i, j in enumerate(arr):
  for k, e in enumerate(j):
    if to_black[k]:
      arr[i,k] = 0
print(arr)

That outputs this:

[[ 0  1  2  3  4  5  6  7  8]
 [ 9 10 11 12 13 14 15 16 17]
 [18 19 20 21 22 23 24 25 26]]
[[ 0  0  2  0  0  5  6  0  8]
 [ 0  0 11  0  0 14 15  0 17]
 [ 0  0 20  0  0 23 24  0 26]]

Now I wonder how one would do it faster in terms of CPU performance?

numpy python
1 Answers

You picked an easy one for numpy.

timr@tims-gram:~/src$ cat x.py
import numpy as np
arr = np.arange(3*9).reshape((3, 9))
print(arr)
to_black = np.random.choice(a=[False, True], size=9)
arr[:,to_black] = 0
print(arr)
timr@tims-gram:~/src$ python x.py
[[ 0  1  2  3  4  5  6  7  8]
 [ 9 10 11 12 13 14 15 16 17]
 [18 19 20 21 22 23 24 25 26]]
[[ 0  0  0  0  4  5  6  0  8]
 [ 9  0  0  0 13 14 15  0 17]
 [18  0  0  0 22 23 24  0 26]]
timr@tims-gram:~/src$
user_18833160
Related