Python: Как найти индексы элементов, удовлетворяющих условиям в каждой строке, и преобразовать их в dict?

#python #arrays #numpy

Вопрос:

Пример:

 import numpy as np
np.random.seed(20211021)
myarray = np.random.randint(0, 5, size=(5, 4))

>>> myarray
array([[2, 3, 0, 1],
       [3, 3, 3, 1],
       [1, 0, 0, 0],
       [3, 2, 4, 0],
       [4, 1, 4, 0]])
 

Здесь я использую argwhere in numpy для поиска индексов элементов, которые больше 0 в каждой строке.

 g0 = np.argwhere(myarray > 0)
>>> g0
array([[0, 0],
       [0, 1],
       [0, 3],
       [1, 0],
       [1, 1],
       [1, 2],
       [1, 3],
       [2, 0],
       [3, 0],
       [3, 1],
       [3, 2],
       [4, 0],
       [4, 1],
       [4, 2]], dtype=int64)
 

Dices g0 — это двумерный массив. Форма индексов, которые я намерен создать, выглядит следующим образом:

 {
    0: [0, 1, 3],
    1: [0, 1, 2, 3],
    2: [0],
    3: [0, 1, 2],
    4: [0, 1, 2]
}
 

Есть ли какой-либо способ, которым g0 можно преобразовать в dict?
(Кроме применения функции к каждой строке myarray , я не нашел эффективного метода)

Комментарии:

1. построение словаря с for помощью цикла не так уж и плохо…

Ответ №1:

np.unique может использоваться с индексами для получения как ключей словаря, так и местоположений, затем используется np.split для разделения массива, затем zip вместе с ключами и массивами для построения словаря из кортежей:

 g0 = np.argwhere(myarray > 0)
keys, locs = np.unique(g0[:, 0], return_index=True)
d = dict(zip(keys, np.split(g0[:, 1], locs[1:])))
 

np.nonzero может быть быстрее, чем np.argwhere в этом случае:

 i, v = np.nonzero(myarray > 0)
keys, locs = np.unique(i, return_index=True)
d = dict(zip(keys, np.split(v, locs[1:])))
 

Однако простое понимание словаря, вероятно, является самым быстрым вариантом для небольших массивов:

 d = {i: np.nonzero(r > 0)[0] for i, r in enumerate(myarray)}
 

Все параметры выдают d :

 {0: array([0, 1, 3]),
 1: array([0, 1, 2, 3]),
 2: array([0]),
 3: array([0, 1, 2]),
 4: array([0, 1, 2])}
 

Настройка и импорт:

 import numpy as np

np.random.seed(20211021)
myarray = np.random.randint(0, 5, size=(5, 4))