基于另一个数组定义的组的数组差异

Differences in an array based on groups defined by another array

我有两个相同大小的数组。一,称它为A,包含一串重复的数字;另一个 B 包含随机数。

import numpy as np

A = np.array([1,1,1,2,2,2,0,0,0,3,3])
B = np.array([1,2,3,6,5,4,7,8,9,10,11])

我需要找到 A 中的组定义的两个极端之间 B 的差异。更具体地说,我需要一个输出 C 例如

C = [2, -2, 2, 1]

其中每一项都是差值 3 - 1、4 - 6、9 - 7 和 11 - 10,即 B 中由 A.

我试着用 itertools.groupby 来隔离第一个数组中的组,但我不清楚如何利用索引来操作第二个数组中的差异。

这很容易实现 pandas' groupby:

A = np.array([1,1,1,2,2,2,0,0,0,3,3])
B = np.array([1,2,3,6,5,4,7,8,9,10,11])

import pandas as pd
pd.Series(B).groupby(A, sort=False).agg(lambda g: g.iloc[-1]-g.iloc[0]).to_numpy()

输出:array([ 2, -2, 2, 1])

使用 itertools.groupby:

from itertools import groupby

[(x:=list(g))[-1][1]-x[0][1] for k, g in groupby(zip(A,B), lambda x: x[0])]

输出:[2, -2, 2, 1]

注意。请注意,如果有不同的non-consecutive组

,这两个解决方案的行为会有所不同

编辑:C 现在的排序方式与问题中的排序方式相同

C = []
_, idx = np.unique(A, return_index=True)
for i in A[np.sort(idx)]:
    bs = B[A==i]
    C.append(bs[-1] - bs[0])

print(C) // [2, -2, 2, 1]

np.unique returns,对于A中的每个唯一值,它第一次出现的索引。

i in A[np.sort(idx)] 按照索引的顺序迭代唯一值。

B[A==i] 在与 A 中的值相同的索引处从 B 中提取值。