使用一个值数组将numpy分成多个子数组

6

我有一组沿着一条线的点:

a = np.array([18, 56, 32, 75, 55, 55])

我有另一个数组,对应于我想要使用它来访问 a 中的信息的索引(它们的长度始终相等)。a 数组和 b 数组都没有排序。

b = np.array([0, 2, 3, 2, 2, 2])

我希望将a分组为多个子数组,以便实现以下功能:
c[0] -> array([18])
c[2] -> array([56, 75, 55, 55])
c[3] -> array([32])

尽管上面的例子很简单,但我将处理数百万个点,因此更喜欢高效的方法。稍后,在程序中可以通过自动方法以这种方式访问任何子点阵也是至关重要的。

1
根据您的目标,这可能会有所帮助:使用NumPy进行向量化分组 - Brad Solomon
1个回答

5
这里有一种方法 -
def groupby(a, b):
    # Get argsort indices, to be used to sort a and b in the next steps
    sidx = b.argsort(kind='mergesort')
    a_sorted = a[sidx]
    b_sorted = b[sidx]

    # Get the group limit indices (start, stop of groups)
    cut_idx = np.flatnonzero(np.r_[True,b_sorted[1:] != b_sorted[:-1],True])

    # Split input array with those start, stop ones
    out = [a_sorted[i:j] for i,j in zip(cut_idx[:-1],cut_idx[1:])]
    return out

一个更简单但效率较低的方法是使用np.split来替换最后几行并获取输出,如下所示 -
out = np.split(a_sorted, np.flatnonzero(b_sorted[1:] != b_sorted[:-1])+1 )

示例运行 -


In [38]: a
Out[38]: array([18, 56, 32, 75, 55, 55])

In [39]: b
Out[39]: array([0, 2, 3, 2, 2, 2])

In [40]: groupby(a, b)
Out[40]: [array([18]), array([56, 75, 55, 55]), array([32])]

为了获取覆盖 b 中所有 ID 范围的子数组 -
def groupby_perID(a, b):
    # Get argsort indices, to be used to sort a and b in the next steps
    sidx = b.argsort(kind='mergesort')
    a_sorted = a[sidx]
    b_sorted = b[sidx]

    # Get the group limit indices (start, stop of groups)
    cut_idx = np.flatnonzero(np.r_[True,b_sorted[1:] != b_sorted[:-1],True])

    # Create cut indices for all unique IDs in b
    n = b_sorted[-1]+2
    cut_idxe = np.full(n, cut_idx[-1], dtype=int)

    insert_idx = b_sorted[cut_idx[:-1]]
    cut_idxe[insert_idx] = cut_idx[:-1]
    cut_idxe = np.minimum.accumulate(cut_idxe[::-1])[::-1]

    # Split input array with those start, stop ones
    out = [a_sorted[i:j] for i,j in zip(cut_idxe[:-1],cut_idxe[1:])]
    return out

样例运行 -

In [241]: a
Out[241]: array([18, 56, 32, 75, 55, 55])

In [242]: b
Out[242]: array([0, 2, 3, 2, 2, 2])

In [243]: groupby_perID(a, b)
Out[243]: [array([18]), array([], dtype=int64), 
           array([56, 75, 55, 55]), array([32])]

1
你为什么在 argsort() 中指定 kind="mergesort" - berkelem
2
@berkelem 为了保持a中元素的顺序不变。使用默认设置时,不能保证保持原有顺序。 - Divakar

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接