Python：如何在每一行中找到满足条件的元素的索引，并将它们转换为字典？

Question

An example:

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]])

Here I use argwhere in numpy to find indices of elements that greater than 0 in each row.

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)

The dices g0 is a two-dimension array. The form of indices that I intend to create is like below:

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

Is there any way in which g0 can be transformed to a dict? (Other than applying function to each row of myarray I hasn't find an efficient method)

Answer 1

np.unique can be used with indexes to get both the dictionary keys and locations, then use np.split to divide the array, then zip together the keys and the arrays to build the dictionary from the tuples:

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 may be faster than np.argwhere in this case:

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

However, a simple dictionary comprehension is likely the fastest option on smaller arrays:

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

All options produce 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])}

---

Setup and imports:

import numpy as np

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