在 numpy 中查找所有匹配值的 3D 索引

Question

我有一个 3D int64 Numpy 数组，它是从 skimage.measure.label 输出的。我需要一个 3D 索引列表，这些索引与我们每个可能的（以前已知的）值相匹配，并由对应于每个值的索引分开。

目前，我们通过以下成语来做到这一点：

    for cur_idx,count in values_counts.items():
        region=labels[:,:,:] == cur_idx
        [dim1_indices,dim2_indices,dim3_indices]= np.nonzero(region)

虽然这段代码可以工作并产生正确的输出，但它很慢，尤其是 np.nonzero 部分，因为我们在一个大数组上调用了 200 多次。我意识到可能有一种更快的方法可以通过 numba 来执行此操作，但我们希望避免添加额外的要求，除非需要。

最终，我们要寻找的是相对有效地对应于每个（非零）值的索引列表。假设我们的值数量100x1000x1000。因此，例如，在由以下创建的数组上：

x = np.zeros((4,4,4))
x[3,3,3] = 1; x[1,0,3] = 2; x[1,2,3] = 2

我们需要一些idx_value dict/array ，例如idx_value_1[2] = 1 idx_value_2[2] = 2, idx_value_3[2] = 3。

Answer 1

我尝试解决与您描述的问题类似的问题，我认为np.argwhere 函数可能是减少运行时间的最佳选择（请参阅文档here）。请参阅下面的代码示例，了解如何根据您在上面确定的约束来使用它。

import numpy as np

x = np.zeros((4,4,4))
x[3,3,3] = 1; x[1,0,3] = 2; x[1,2,3] = 3

# Instantiate dictionary/array to store indices
idx_value = {}

# Get indices equal to 2
idx_value[3] = np.argwhere(x == 3)
idx_value[2] = np.argwhere(x == 2)
idx_value[1] = np.argwhere(x == 1)

# Display idx_value - consistent with indices we set before
>>> idx_value
{3: array([[1, 2, 3]]), 2: array([[1, 0, 3]]), 1: array([[3, 3, 3]])}

对于第一个用例，我认为您仍然需要使用 for 循环来迭代您正在搜索的值，但可以这样做：

# Instantiate dictionary/array
idx_value = {}

# Now loop by incrementally adding key/value pairs
for cur_idx,count in values_counts.items():
        idx_value[cur_idx] = np.argwhere(labels)

注意：这会逐步创建一个字典，其中每个键都是要搜索的 idx，每个值都是 np.array 形状为 (N_matches, 3) 的对象。