python - 返回一维 NumPy 数组中唯一值索引的代码的说明

Question

我在网上找到了这段代码，但由于我不精通 Python，所以很难理解它的每个部分在做什么。

The following routine takes an array as input and returns a dictionary that maps each unique value to its indices

def partition(array):
  return {i: (array == i).nonzero()[0] for i in np.unique(array)}

Answer 1

追踪每个部分，这应该不言自明。内联评论。

In [304]: array = np.array([1, 1, 2, 3, 2, 1, 2, 3])

In [305]: np.unique(array)            # unique values in `array`
Out[305]: array([1, 2, 3])

In [306]: array == 1                  # retrieve a boolean mask where elements are equal to 1
Out[306]: array([ True,  True, False, False, False,  True, False, False])

In [307]: (array == 1).nonzero()[0]   # get the `True` indices for the operation above
Out[307]: array([0, 1, 5])

总结；该代码正在创建 <unique_value : all indices of unique_value in array> 的映射-

In [308]: {i: (array == i).nonzero()[0] for i in np.unique(array)}
Out[308]: {1: array([0, 1, 5]), 2: array([2, 4, 6]), 3: array([3, 7])}

这里是更易读的版本 -

In [313]: mapping = {}
     ...: for i in np.unique(array):
     ...:     mapping[i] = np.where(array == i)[0] 
     ...:     

In [314]: mapping
Out[314]: {1: array([0, 1, 5]), 2: array([2, 4, 6]), 3: array([3, 7])}