python - 在 numpy 数组中找到人口最多的 "slice"

Question

我有一个形状为 [x_len,y_len,z_len] 的 numpy 掩码。我希望找到使 np.count_nonzero(mask[:,:,z]) 最大化的 z。

我天真的解决方案:

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

但我正在寻找更快和/或更漂亮的东西。

Answer 1

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

产生最大非零元素的 foo 的 z-index。

---

为了比较此解决方案与@mcsoini 的解决方案和另一个新颖的解决方案:

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)

最后一个解决方案花费的时间是其他两个解决方案的一半。我们可以负担得起这个技巧，因为掩码实际上是 0 和 1 值的矩阵。如果有其他值，它将不起作用。

---

进一步评论:

如果 foo 是 bool 类型(掩码应该是)，那么所有这些方法似乎都需要完全相同的时间(在误差范围内)，表明，也许在幕后， bool 值的 count_nonzero 与 sum 非常相似？不过，我不知道，如果有一些见解会很好。