Извлечение необходимых строк из массива numpy

Question

У меня есть заданный массив numpy следующим образом.

import numpy as np

data = np.array([[4,6,8,9,3,2,4,4,1], # no of 0s == 0
                  [4,6,8,9,3,0,0,4,0], # no of 0s == 3
                  [4,6,0,9,0,2,0,4,0], # no of 0s == 4
                  [4,6,8,0,3,0,0,0,0], # no of 0s == 5
                  [4,6,8,9,3,2,0,4,0]]) # no of 0s == 2

Из данного массива данных, я должен извлечь 3 строки, которые содержат наименьшее 0. Итак, ожидаются 1-й, последний и второй ряды.

res = np.array([[4,6,8,9,3,2,4,4,1], # no of 0s == 0
                      [4,6,8,9,3,0,0,4,0], # no of 0s == 3                      
                      [4,6,8,9,3,2,0,4,0]]) # no of 0s == 2

Как я могу это сделать, ребята?

Answer 1

Сумма по вашему условию и разделу.

n = 3
c = (data == 0).sum(1)
mn = np.argpartition(c, n)[:n]
data[mn]

array([[4, 6, 8, 9, 3, 2, 4, 4, 1],
       [4, 6, 8, 9, 3, 2, 0, 4, 0],
       [4, 6, 8, 9, 3, 0, 0, 4, 0]])

Если вам нужны строки, отсортированные по исходному значению индекса, а не по количеству нулей, замените последнюю строку на:

data[np.sort(mn)]