重复选择函数微小变化的快速数值计算

Question

我想用numpy来解决一个与纽比。重复功能。我不知道如何使用我熟悉的numpy函数来解决这个问题，所以我正在寻找帮助，看看这是否可以用numpy来完成。我的数组很大（>；1e6个元素），高性能至关重要，因此我无法承受python for循环的性能损失。在

最小示例

我有一个length-num pts排序整数数组objID，它存储（可能重复）对象标识符。在

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

使用cd1}的唯一索引来确定它们的唯一索引纽比。独一无二. 在

^{2}$

我还有一个length-num惟一的objIDs数组occupations，它指定我要从objID中选择unique_objIDs的每个条目多少次。在

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

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

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的第一个元素。该值为零，它告诉我们，我们不想选择存储unique_objIDs[0]=0的objID的任何索引，因此所有这些索引都被排除在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中的前2个出现的索引。这就是desired_array_of_indices的最后两个条目是6和7的原因。在

区别于np.重复在

注意这个计算和numpy.repeat之间的细微差别。对于numpy.repeat，返回的索引属于唯一项数组unique_objIDs。这里我需要objID的索引，我还需要为重复条目的情况选择连续的索引。可以假定occupations的每个条目小于或等于相应条目在objID中出现的总次数，因此不存在索引错误的危险。在

有人知道如何用可用的向量化Numpy函数（可能是一些集合）来表述这个问题吗？在

Answer 1

另一个建议是用return_counts代替return_index：

unique_objIDs, objID_counts = np.unique(objID, return_counts=True)
num_unique_objIDs = len(unique_objIDs)

yesno = np.tile([True, False], num_unique_objIDs)
amounts = np.c_[occupations, objID_counts-occupations].ravel()
desired_array_of_indices = np.flatnonzero(np.repeat(yesno, amounts))

Answer 2

有一个办法。在

首先，您的示例代码：

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的更有效方法是：

^{pr2}$

此操作是O（n），而不是unique所需的O（n*log（n））。然后可以使用

unique_objIDs = objID[idx_unique_objIDs]

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

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

现在找到所需的指数。结果在Out[107]行中：

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

仔细观察：

csum是occupations的累计和，n是{}的总和：

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

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

csum可以解释为与每个职业相关联的指数范围的结束。然后csum - occupations保存范围开始的索引：

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

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

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

如果从np.arange(n)中减去这个值，那么对于每个职业k，我们有一个从0到occupation[k]-1的范围串联在一个数组中：

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

那不是理想的结果。我们必须添加（repeated）idx_unique_objIDs，使值成为数组objID的索引：

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

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

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