给定numpy.where()的输出,如何获取索引列表的列表



我正在为零元素搜索一个相当大的矩阵,然后想要循环遍历行和每行内部,总体上我找到了一个零的索引。numpy.where(M==0)得到(np.array, np.array)

我如何以最有效的方式格式化我需要的结果?

# Output of numpy.where(M==0):
(x,y) = (array([0, 0, 1, 1, 2]), array([2, 3, 1, 2, 1]))
# Output I want, each value of y in the list corresponding to the x value:
wanted = [[2,3],[1,2],[1],[]]

我考虑过从输出中构造一个稀疏矩阵(因为M是密集的,M==0应该是稀疏的),但不知道如何迭代那里的行。

zeros = scipy.sparse.coo_matrix((np.ones((len(x),)), (x,y)), shape=M.shape)
print(zeros.nonzero()) # is the same output as np.where, so it does not help me

我想我可以切成每一行并.nonzero()它,但我希望有更有效的东西

编辑:或者,{x_value:[list of corresponding y_values]}形式的字典也可以工作

我发现collections.defaultdict允许我将元组迭代到我想要的字典中。

indices = np.argwhere(M==0)
zeros = defaultdict(list)
for coords in indices:
zeros[coords[0]].append(coords[1])
print(zeros)
# Output:
defaultdict(<class 'list'>, {0: [1], 1: [1, 2], 2: [0, 3]})

最新更新