提取包含特定值的數組索引

Question

我有一個形狀如下的數組： （31777，44，44）

它的sum為31777.0，並且每個(44,44)網格在某些坐標處僅包含一個1.0條目。

我的目標是獲得一個形狀為（31777，2）的數組，其中2是網格中1.0條目的坐標。 這樣做的麻煩方法是循環，但是我正在尋找一種優雅的高性能解決方案。

Answer 1

您可以使用numpy.where在多維數組中查找元素的索引，例如：

import numpy as np
arr = np.array([[[1,0],[0,0]], [[0,0],[1,0]], [[0,0], [0,1]]])
arr.shape
# (3, 2, 2)

np.array([coord for coord in zip(*np.where(arr == 1)[1:])])
# array([[0, 0],
#        [1, 0],
#        [1, 1]]) this is an array of shape (3,2) which contains coordinates of 1.0 entry.

Answer 2

第一個答案給了我Google材料，以查找我認為是最佳的：

np.vstack(arr.nonzero()[1:][::-1]).T