Here are 3 ways to be quicker at calculating row_sum, col_sum and ones_count.
Baseline
For testing I use this array
import numpy as np
import numba as nb
np.random.seed(1)
n = 10**4
array = np.random.randint(0,2,(n,n))
Now your exact code takes 20.3 s ± 397 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) on my machine.
The Lazy One Liner Numpy Version:
%timeit np.stack(np.where(array)).sum(axis=1),array.sum() takes 1.13 s ± 12.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) on my machine.
Here np.stack(np.where(array)).sum(axis=1) is what you call row_sum and col_sum and array.sum() gives your ones_count
Avoid Looping Threw Twice
You can use your exact code with numba.jit
@nb.njit
def test():
row_sum = 0
col_sum = 0
ones_count = 0
for row_count, row in enumerate(array):
for col_count, col in enumerate(row):
if col == 1:
row_sum += row_count
col_sum += col_count
ones_count += 1
return row_sum,col_sum,ones_count
%timeit test()
This is a bit faster. It takes 50 ms ± 614 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) on my machine. But defininitly not worth the effort.
Multicore Version
A slight modification of your code can run multithreaded with numba
@nb.njit(parallel=True)
def test2():
row_sum = 0
col_sum = 0
ones_count = 0
for row_count in nb.prange(len(array)):
row = array[row_count]
for col_count, col in enumerate(row):
if col == 1:
row_sum += row_count
col_sum += col_count
ones_count += 1
return row_sum,col_sum,ones_count
%timeit test2()
Now this does give a little speed up compared to the lazy numpy version. It takes 13.3 ms ± 2.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) on my 10 core machine. Though it's not using all 10 cores.
Notice that you have to be careful when modifying things in parallel. You could create a race condition. And this is not the case here only because numba takes counter measures for this specific case.
Further Optimisations
As pointed out by Jérôme Richard in a comment. The last version can be optimised by using uint8 to sture the instead of int64 which is the default. Just call .astype(np.uint8) on the array. Then it takes 9.38 ms ± 935 µs per loop (mean ± std. dev. of 7 runs, 1 loop each) on my machine.
