python - 快速Numpy计算重复和选择功能的微小变化

我想使用numpy来解决一个与numpy.repeat函数解决的问题非常相似但不完全相同的问题。我不知道如何使用我熟悉的任何numpy函数来解决这个问题，所以我正在寻求帮助，看看是否可以用numpy来解决这个问题。我的数组很大（>1e6元素），高性能是关键，因此我无法承受python for循环的性能影响。
最小示例
我有一个length-num-pts排序整数数组objID存储（可能重复）对象标识符。

objID = np.array([0, 0, 5, 5, 5, 7, 8, 8])

我使用numpy.unique确定objID的唯一条目，以及它们在objID中的外观索引。

unique_objIDs, idx_unique_objIDs = np.unique(objID, return_index=True)
num_unique_objIDs = len(unique_objIDs)

我还有一个length-num_unique_objIDs数组occupations指定要从unique_objIDs中选择objID的每个条目的次数。

occupations = np.array([0, 2, 1, 2])

我想根据objID确定可以用来检索occupations元素的索引数组。下面我举一个具体的例子。

desired_array_of_indices = np.array([2, 3, 5, 6, 7])

数组desired_array_of_indices是我想用numpy计算的。desired_array_of_indices的条目计算如下。
desired_array_of_indices的明确解释
occupations数组的Element-i指定选择unique_objID[i]的次数。desired_array_of_indices数组存储这些选择的objID索引。对于多次选择的objID值，选择连续的objID索引，这样就不会重复存储在desired_array_of_indices中的索引。
对于具体性，考虑occupations的第一个元素。该值为零，告诉我们不想选择存储objID的任何unique_objIDs[0]=0索引，因此所有这些索引都不在desired_array_of_indices中。
occupations的下一个元素是2，告诉我们要选择unique_objIDs[1]=5中objID的前两个外观的索引。这就是为什么desired_array_of_indices的前两个条目是2和3。
occupations的下一个元素是1，它告诉我们要选择unique_objIDs[2]=7中第一个出现objID的索引。所以desired_array_of_indices的下一个条目是5。
occupations的最后一个元素是2，告诉我们要选择unique_objIDs[3]=8中objID的前两个外观的索引。这就是为什么desired_array_of_indices的最后两个条目是6和7。
与np.repeat的区别
注意这个计算和numpy.repeat之间的细微差别。对于numpy.repeat，返回的索引属于唯一条目的数组unique_objIDs。这里我需要objID的索引，并且我还需要为重复输入的情况选择连续索引。可以假定occupations的每个条目小于或等于对应的条目出现在objID中的总数，因此不会有索引错误的危险。
有人看到如何用（可能是一些）可用的矢量化Numpy函数来描述这个问题吗？

最佳答案

有一个办法。
首先，您的示例代码：

In [102]: objID = np.array([0, 0, 5, 5, 5, 7, 8, 8])

In [103]: unique_objIDs, idx_unique_objIDs = np.unique(objID, return_index=True)

[[注：unique()对其参数排序。您知道您的输入已经排序，因此获得idx_unique_objIDs的更有效方法是：

idx_unique_objIDs = np.concatenate(([0], np.nonzero(np.diff(objID))[0] + 1))

此操作是O（n），而不是unique所需的O（n*log（n））。那你就可以用

unique_objIDs = objID[idx_unique_objIDs]

如果您需要唯一对象ID的数组。]]

In [104]: occupations = np.array([0, 2, 1, 2])

现在找到所需的索引。结果如下：

In [105]: csum = occupations.cumsum()

In [106]: n = csum[-1]

In [107]: np.arange(n) + np.repeat(idx_unique_objIDs - csum + occupations, occupations)
Out[107]: array([2, 3, 5, 6, 7])

仔细看：
Out[107]是csum的累积和，occupations是n的总和：

In [114]: csum
Out[114]: array([0, 2, 3, 5])

In [115]: n
Out[115]: 5

occupations可以解释为与每个职业相关的指数范围（pythonic“end”，即）的结束指数。然后csum保存范围开始的索引：

In [116]: csum - occupations
Out[116]: array([0, 0, 2, 3])

根据csum - occupations中的值重复这些起始索引：

In [117]: np.repeat(csum - occupations, occupations)
Out[117]: array([0, 0, 2, 3, 3])

如果从occupations中减去该值，则对于每个职业np.arange(n)，我们将0到k之间的范围串联在一个数组中：

In [118]: np.arange(n) - np.repeat(csum - occupations, occupations)
Out[118]: array([0, 1, 0, 0, 1])

那不是很理想的结果。我们必须添加（重复的）occupation[k]-1，以便这些值成为数组中的索引：

In [119]: np.arange(n) - np.repeat(csum - occupations, occupations) + np.repeat(idx_unique_objIDs, occupations)
Out[119]: array([2, 3, 5, 6, 7])

现在将这两个idx_unique_objIDs调用组合起来得到最终表达式：

In [120]: np.arange(n) + np.repeat(idx_unique_objIDs - csum + occupations, occupations)
Out[120]: array([2, 3, 5, 6, 7])

关于python - 快速Numpy计算重复和选择功能的微小变化，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/37632203/