请求解析返回一维NumPy数组唯一值索引的Python代码逻辑

嘿，我来帮你一步步拆解这段Python代码，它依赖NumPy库实现了一个实用功能——把输入数组里的每个唯一值，对应到它在原数组中所有出现过的索引位置。

首先先补全完整可运行的代码（原代码漏了NumPy导入，运行前必须加）：

import numpy as np

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

接下来逐个部分拆解：

• 获取数组的唯一值：np.unique(array)是NumPy的内置函数，它会返回输入数组中所有不重复的值，而且默认是按升序排列的。比如输入np.array([2,1,2,3,1])，这部分会返回array([1,2,3])。
• 元素级比较生成布尔数组：array == i是NumPy的向量化比较操作，会生成一个和原数组形状完全一致的布尔数组——原数组中等于i的位置是True，其余位置是False。还是拿上面的例子，当i=2时，这一步会得到array([True, False, True, False, False])。
• 提取True值的索引：.nonzero()方法会找出布尔数组中所有True对应的位置，返回一个元组（如果是一维数组，元组里只有一个元素）。我们用[0]取出这个元素，就得到了所有符合条件的索引组成的一维数组。
• 字典推导式构建结果：整个{i: ... for i in ...}是Python的字典推导式，把每个唯一值i作为字典的键，对应的索引数组作为值，最终返回这个映射字典。

举个实际例子更直观 #

假设我们输入：

arr = np.array([5, 3, 5, 7, 3, 3])
print(partition(arr))

运行后会输出：

{3: array([1, 4, 5]), 5: array([0, 2]), 7: array([3])}

注意事项 #

• 这段代码必须依赖NumPy，所以运行前一定要先导入import numpy as np，否则会报错。
• 输入的array最好是NumPy数组，如果传入普通Python列表，array == i会变成判断整个列表是否等于i（而不是元素级比较），导致结果错误。如果是列表，建议先转成NumPy数组再传入，比如partition(np.array([1,2,1]))。

内容的提问来源于stack exchange，提问作者SomebodyOnEarth