Find the most populated "slice" in a numpy array

Question

I have a numpy mask with shape [x_len,y_len,z_len]. I wish to find the z such that np.count_nonzero(mask[:,:,z]) is maximised.

My naive solution:

best_z = -1
best_score = -1
for z in range(mask.shape[2]):
    n_nonzero = np.count_nonzero(mask[:,:,z])

    if n_nonzero > best_score:
        best_score = n_nonzero
        best_z = z

But I'm looking for something faster and/or prettier.

Answer 1

np.argmax(np.count_nonzero(foo, axis=(0, 1)))

yields the z-index of foo for which there are maximum non-zero elements.

---

For a comparison of this solution, with @mcsoini's solution and another novel one:

foo = np.random.randint(0, 2, size=(100, 100, 200))

# this solution
i> %timeit np.argmax(np.count_nonzero(foo, axis=(0, 1)))
o> 1.58 ms ± 43.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# @mcsoini's solution
i> %timeit np.argmax(np.count_nonzero(foo.reshape(-1, foo.shape[-1]), axis=0))
o> 1.64 ms ± 18.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# a trick solution
i> %timeit np.argmax(np.sum(foo, axis = (0, 1)))
o> 709 µs ± 4.87 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

The last solution takes half the time of the other two solutions. We can afford this trick since a mask is effectively a matrix of 0 and 1 values. It won't work if there are other values.

---

Further comments:

It seems like all these methods take exactly the same time (within margin of error) if foo is of type bool (which a mask is expected to be), indicating, perhaps under the hood, count_nonzero for boolean values is very similar to sum? I don't know, though, it would be nice to have some insight.