假设我们有一个秩为2的数组
一个天真的方法涉及布尔索引,这将导致
(Note: 我随后需要这些索引以进行写入访问,也就是替换数组中存储的值。但在这些操作之间,我确实需要这个二维结构。)
a
,其中包含n
个整数值,这些值在 {0,1,2,...,m}
中。现在对于每个整数,我想找到具有该值的 a
条目的索引(在以下示例中称为 index_i、index_j
)。 (因此,我要寻找类似于 np.unique(..., return_index=True)
的东西,但是针对 2d 数组,并且可以返回每个唯一值的所有索引。)一个天真的方法涉及布尔索引,这将导致
O(m*n)
的操作(见下文),但我只想进行 O(n)
操作。虽然我找到了一种解决方案,但我觉得应该有一个内置方法或至少一个简化这个过程的东西 - 或者至少可以消除这些丑陋的循环:import numpy as np
a = np.array([[0,0,1],[0,2,1],[2,2,1]])
m = a.max()
#"naive" in O(n*m)
i,j = np.mgrid[range(a.shape[0]), range(a.shape[1])]
index_i = [[] for _ in range(m+1)]
index_j = [[] for _ in range(m+1)]
for k in range(m+1):
index_i[k] = i[a==k]
index_j[k] = j[a==k]
#all the zeros:
print(a[index_i[0], index_j[0]])
#all the ones:
print(a[index_i[1], index_j[1]])
#all the twos:
print(a[index_i[2], index_j[2]])
#"sophisticated" in O(n)
index_i = [[] for _ in range(m+1)]
index_j = [[] for _ in range(m+1)]
for i in range(a.shape[0]):
for j in range(a.shape[1]):
index_i[a[i,j]].append(i)
index_j[a[i,j]].append(j)
#all the zeros:
print(a[index_i[0], index_j[0]])
#all the ones:
print(a[index_i[1], index_j[1]])
#all the twos:
print(a[index_i[2], index_j[2]])
(Note: 我随后需要这些索引以进行写入访问,也就是替换数组中存储的值。但在这些操作之间,我确实需要这个二维结构。)
numpy.argwhere
?类似于[np.argwhere(a == x) for x in np.unique(a)]
这样的代码? - Chris AdamsO(n*m)
次操作,因为你将每个值x
与整个输入数组进行比较。 - flawrindex_i, index_j
一样。) - flawr