Find the most populated "slice" in a numpy array-CodePudding

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.

CodePudding user response：

I guess this is what you need:

best_z = np.argmax(np.count_nonzero(mask, axis=-1))

CodePudding user response：

You are looking for the index along the z-axis corresponding to the array's slice with the largest number of non-zero elements. With the example data

np.random.seed(3)
mask = np.random.randn(2, 3, 4)
mask = np.where(mask < 0, 0, mask)
print(mask)

[[[1.78862847 0.43650985 0.09649747 0.        ]
  [0.         0.         0.         0.        ]
  [0.         0.         0.         0.88462238]]

 [[0.88131804 1.70957306 0.05003364 0.        ]
  [0.         0.         0.98236743 0.        ]
  [0.         0.         1.48614836 0.23671627]]]

we can first reshape the array mask.reshape(-1, mask.shape[-1]) in order to reduce the dimensions 0 and 1 to a single dimension. Then we count the number of non-zeros along this new first dimension p.count_nonzero(..., axis=0), and finally we can find the indices along z where those counts are maximum (np.argmax):

np.argmax(np.count_nonzero(mask.reshape(-1, mask.shape[-1]), axis=0))

Result: 2

CodePudding user response：

I came up with this:

unique, counts = np.unique(np.where(mask)[2], return_counts=True)
best_z = unique[np.argmax(counts)]

Although I expect dankal and mcsoini's answers are both faster.