Reputation: 303
I'm writing a program in which there is a numpy array a
whose elements can take three possible values: -1, 0 or 1. I am trying to multiply some of its elements by a number c
. The idea is to obtain this behaviour:
for i,el in enumerate(a):
if el == b:
a[i] *= c
I came up with a solution that does not require any loops and works a couple of orders of magnitude faster than the previous one, this is the code I used to test them:
# Long array with random integers between -1 and 1
a = np.random.choice(3,1000000) - 1
a1 = a.copy()
a2 = a.copy()
# Reference values for b and c
b = 1
c = 10
# Solution with loop
t0 = time.time()
for i,el in enumerate(a1):
if el == b:
a1[i] *= c
t1 = time.time()
# Solution without loop
a2 = a2*((a2 == b)*c + (a2 != b))
t2 = time.time()
print("No loop: %f s"%(t1 - t0))
print("Loop: %f s"%(t2 - t1))
Although it seems to be working fine I'm not really happy with multiplying integers by booleans, but I don't know if I should, so I would appreciate if anyone could tell me a bit more about what is Numpy doing and/or if there is a better way to do this that I am not considering.
Thanks in advance!
Upvotes: 2
Views: 1361
Reputation: 231425
Some comparative timings:
In [66]: %%timeit a2=a.copy()
...: a2*((a2==b)*10 + (a2!=b))
14.4 ms ± 36.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [67]: %%timeit a2=a.copy()
...: a2[a2==b] *= 10
1.96 ms ± 75 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [68]: %%timeit a2=a.copy()
...: a2[a2==b] = a2[a2==b]*10
3.28 ms ± 5.63 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [69]: %%timeit a2=a.copy()
...: np.multiply(a2, 10, where=a2==b, out=a2)
1.63 ms ± 3.38 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
The fastest ones only do one a2==b
test. The multiply
with where
parameter is fastest, but also bit harder to understand.
And to verify that the fastest produces the same thing:
In [73]: a2=a.copy();a2=a2*((a2==b)*10 + (a2!=b))
In [74]: a3=a.copy();np.multiply(a3, 10, where=a3==b, out=a3);
In [75]: np.allclose(a2,a3)
Out[75]: True
Upvotes: 0
Reputation: 114881
NumPy will cast the bool type to the integer type, with False
and True
converted to 0 and 1 respectively. This casting is safe, so don't worry, be happy.
In [8]: np.can_cast(np.bool8, np.intc)
Out[8]: True
If you prefer to be explicit, you could do that casting yourself by replacing (a2 == b)
with (a2 == b).astype(int)
, but that is not necessary.
Upvotes: 2