Python:如何找到每行中满足条件的元素的索引，并将其转换为字典?

Question

示例:

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

这里我使用numpy中的argwhere来查找每行中大于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)

器件g0是一个二维阵列。我打算创建的索引的形式如下:

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

是否有办法将g0转化为字典?(除了对myarray的每一行应用函数外，我还没有找到有效的方法)

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

---

Setup and imports:

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