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

Question

I found this snippet of code online and am having difficulty in understanding what each part of it is doing as I'm not proficient in 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

Trace each part out, this should speak for itself. Comments inlined.

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])

In summary; the code is creating a mapping of <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])}

And here's the slightly more readable version -

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])}

Answer 2

• array == i Return a boolean array of True whenever the value is equal to i and False otherwise.
  
• nonzero() Return the indices of the elements that are non-zero(not False). https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.nonzero.html
  
• nonzero()[0] Return the first index where array[index] = i.
  
• for i in np.unique(array) Iterate over all the unique values of array or in other words do the logic foreach value of unique value of the array.

Answer 3

consider also the following Pandas solution:

import pandas as pd

In [165]: s = pd.Series(array)

In [166]: d = s.groupby(s).groups

In [167]: d
Out[167]:
{1: Int64Index([0, 1, 5], dtype='int64'),
 2: Int64Index([2, 4, 6], dtype='int64'),
 3: Int64Index([3, 7], dtype='int64')}

PS pandas.Int64Index - supports all methods and indexing like a regular 1D numpy array

it can be easily converted to Numpy array:

In [168]: {k:v.values for k,v in s.groupby(s).groups.items()}
Out[168]:
{1: array([0, 1, 5], dtype=int64),
 2: array([2, 4, 6], dtype=int64),
 3: array([3, 7], dtype=int64)}