在数组中计算所选差异的高效方法

Question

我有一个模拟脚本的输出，其中包含两个数组，一个包含ID，另一个包含时间，例如：

ids = np.array([2, 0, 1, 0, 1, 1, 2])
times = np.array([.1, .3, .3, .5, .6, 1.2, 1.3])

这些数组始终具有相同的大小。现在我需要计算times的差异，但仅针对具有相同ids的时间。当然，我可以简单地循环遍历不同的ids并执行操作。

for id in np.unique(ids):
    diffs = np.diff(times[ids==id])
    print diffs
    # do stuff with diffs

然而，这种方法效率相当低下，而且两个数组可能非常大。是否有更好的方法可以更有效地完成这项任务？

- obachtos

Answer 1

您可以使用 array.argsort() 并忽略与 ids 变化相对应的值:

>>> id_ind = ids.argsort(kind='mergesort')
>>> times_diffs = np.diff(times[id_ind])
array([ 0.2, -0.2,  0.3,  0.6, -1.1,  1.2])

为了确定需要丢弃哪些值，您可以使用计数器（from collections import Counter）来计算每个ID出现的次数。

或者只需对ID进行排序，并查看其差异不为零的位置：这些是ID更改的索引和时间差不相关的位置：

times_diffs[np.diff(ids[id_ind]) == 0] # ids[id_ind] being the sorted indices sequence

最后，您可以使用np.split和np.where来拆分此数组：

np.split(times_diffs, np.where(np.diff(ids[id_ind]) != 0)[0])

正如您在评论中提到的，argsort() 的默认算法（快速排序）可能不会保留相等元素之间的顺序，因此必须使用 argsort(kind='mergesort') 选项。

- P. Camilleri

---

当你已经有用于排序数组的索引，即ids[id_ind]时，使用sorted(ids)是否有特定原因？ - obachtos

@obachtos 不，那只是懒惰。我会修复它的。 - P. Camilleri

还有一点需要注意：argsort()使用标准算法quicksort可能会破坏时间顺序。最好使用稳定的mergesort，即argsort(kind='mergesort')。 - obachtos

@obachtos 很好的评论。将来请将它作为评论，让我来编辑我的答案：如果您尝试自行编辑，审核人员会拒绝它，因为“此编辑偏离了帖子的原始意图。即使必须进行重大更改的编辑也应该努力保留帖子所有者的目标。” - P. Camilleri

Answer 2

假设你使用ids来调用np.argsort：

inds = np.argsort(ids, kind='mergesort')
>>> array([1, 3, 2, 4, 5, 0, 6])

现在按照 np.diff 的结果对 times 进行排序，并在前面加上一个 nan:

diffs = np.concatenate(([np.nan], np.diff(times[inds])))
>>> diffs 
array([ nan,  0.2, -0.2,  0.3,  0.6, -1.1,  1.2])

这些差异是正确的，除了边界。让我们计算一下这些边界。

boundaries = np.concatenate(([False], ids[inds][1: ] == ids[inds][: -1]))
>>> boundaries
array([False,  True, False,  True,  True, False,  True], dtype=bool)

现在我们只需要这样做

diffs[~boundaries] = np.nan

---

让我们来看看我们得到了什么：

>>> ids[inds]
array([0, 0, 1, 1, 1, 2, 2])

>>> times[inds]
array([ 0.3,  0.5,  0.3,  0.6,  1.2,  0.1,  1.3])

>>> diffs
array([ nan,  0.2,  nan,  0.3,  0.6,  nan,  1.2])

- Ami Tavory

Answer 3

我会翻译中文。下面是需要翻译的内容：

我正在添加另一个答案，尽管这些事情在numpy中是可能的，但我认为更高级别的pandas对它们来说更自然。

在pandas中，您可以在创建DataFrame后一步完成此操作：

df = pd.DataFrame({'ids': ids, 'times': times})

df['diffs'] = df.groupby(df.ids).transform(pd.Series.diff)

This gives:

>>> df
   ids  times  diffs
0    2    0.1    NaN
1    0    0.3    NaN
2    1    0.3    NaN
3    0    0.5    0.2
4    1    0.6    0.3
5    1    1.2    0.6
6    2    1.3    1.2

- Ami Tavory

---

这是一个很好的答案。让我补充一下，有时候你的数据框有多列。在需要时包含这些列是一个好主意。df['diffs'] = df.groupby(['ids'])['times'].transform(pd.Series.diff) - Shane S

Answer 4

numpy_indexed 包（免责声明：我是它的作者）包含了这种分组操作的高效灵活功能：

import numpy_indexed as npi
unique_ids, diffed_time_groups = npi.group_by(keys=ids, values=times, reduction=np.diff)

与pandas不同，它不需要专门的数据结构来执行这种相当基本的操作。

- Eelco Hoogendoorn

---

通常情况下，当有人推广自己的库时，习惯上会添加免责声明，说明他/她是作者。 - Ami Tavory

啊，是的；我习惯这样做，但我忘了；谢谢。 - Eelco Hoogendoorn

祝你的软件包好运。 - Ami Tavory