我想使用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数组的元素i指定将选择unique_objID[i]的次数。 desired_array_of_indices数组存储这些选择的objID索引。对于多次选择的objID值，选择连续索引objID，以便不重复存储在desired_array_of_indices中的索引。
具体而言，请考虑occupations的第一个要素。该值为零，告诉我们我们不想选择存储objID的{{1}}的任何索引，因此所有这些索引都不在unique_objIDs[0]=0之内。
desired_array_of_indices的下一个元素是2，告诉我们要在occupations中选择unique_objIDs[1]=5的前两个出现的索引。这就是objID的前两个条目是2和3的原因。
desired_array_of_indices的下一个元素是1，告诉我们要在occupations中选择unique_objIDs[2]=7的首次出现的索引。所以objID的下一个条目是5.
desired_array_of_indices的最后一个元素是2，告诉我们要在occupations中选择unique_objIDs[3]=8的前两个出现的索引。这就是objID的最后两个条目是6和7的原因。
请注意此计算与desired_array_of_indices之间的细微差别。对于numpy.repeat，返回的索引与唯一条目数组numpy.repeat相关。这里我需要unique_objIDs的索引，我还需要为重复输入的情况选择连续的索引。可以假定objID的每个条目小于或等于相应条目在occupations中出现的总次数，因此不存在索引错误的危险。
有没有人看到如何根据可能的矢量化Numpy函数（可能是一些集合）来制定这个问题？
答案 0 :(得分：4)
这是单向的。
首先，您的示例代码：
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])
现在找到所需的指数。结果在第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是occupations的总和：
In [114]: csum
Out[114]: array([0, 2, 3, 5])
In [115]: n
Out[115]: 5
csum可以解释为与每个职业相关的索引范围（pythonic＆＃34; end＆＃34;，即）的 end 的索引。然后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])
这不是理想的结果。我们必须添加（重复）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])
答案 1 :(得分：2)
另一个建议是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))