python - 在python中，是否有一种有效的方法来将一个数组与映射到另一个数组的元素分开？

Question

假设我有一个任意数组 np.array([1,2,3,4,5,6])另一个数组将数组中的特定元素映射到一个组，np.array(['a','b', 'a','c','c', 'b'])现在我想根据第二个数组中给出的标签/组将它们分成三个不同的数组，以便它们是 a,b,c = narray([1,3]), narray([2,6]), narray([4,5]) .是一个简单的forloop要走的路还是我在这里缺少一些有效的方法？

Answer 1

当您编写高效时，我假设您在这里想要的实际上是快速的。
我将尝试简要讨论渐近效率。
在这种情况下，我们指的是 N作为输入大小和 K作为唯一值的数量。
我的方法解决方案是使用 np.argsort() 的组合和定制的 groupby_np()专门针对 NumPy 输入进行了优化:

import numpy as np


def groupby_np(arr, both=True):
    n = len(arr)
    extrema = np.nonzero(arr[:-1] != arr[1:])[0] + 1
    if both:
        last_i = 0
        for i in extrema:
            yield last_i, i
            last_i = i
        yield last_i, n
    else:
        yield 0
        yield from extrema
        yield n


def labeling_groupby_np(values, labels):
    slicing = labels.argsort()
    sorted_labels = labels[slicing]
    sorted_values = values[slicing]
    del slicing
    result = {}
    for i, j in groupby_np(sorted_labels, True):
        result[sorted_labels[i]] = sorted_values[i:j]
    return result

这很复杂 O(N log N + K) . N log N来自排序步骤和 K来自最后一个循环。
有趣的部分是 N -依赖和 K -依赖步骤很快，因为 N -dependent部分在低层执行， K -依赖部分是 O(1)而且也很快。

类似以下的解决方案(非常类似于 @theEpsilon 答案):

import numpy as np


def labeling_loop(values, labels):
    labeled = {}
    for x, l in zip(values, labels):
        if l not in labeled:
            labeled[l] = [x]
        else:
            labeled[l].append(x)
    return {k: np.array(v) for k, v in labeled.items()}

使用两个循环并具有 O(N + K) .我认为您无法轻松避免第二个循环(没有显着的速度损失)。至于第一个循环，它是在 Python 中执行的，它本身会带来显着的速度损失。

另一种可能性是使用 np.unique() 这将主循环带到了较低的级别。然而，这带来了其他挑战，因为一旦提取了唯一值，没有一些 NumPy advanced indexing 就没有有效的方法来提取信息来构建您想要的数组。 , 即 O(N) .这些解决方案的整体复杂性是 O(K * N) ，但是因为 NumPy 高级索引是在较低级别完成的，所以这可以找到相对较快的解决方案，尽管比替代方案具有更差的渐进复杂性。
可能的实现包括(类似于 @AjayVerma's 和 @AKX's 答案):

import numpy as np


def labeling_unique_bool(values, labels):
    return {l: values[l == labels] for l in np.unique(labels)}

import numpy as np


def labeling_unique_nonzero(values, labels):
    return {l: values[np.nonzero(l == labels)] for l in np.unique(labels)}

此外，可以考虑预先排序步骤，然后通过避免 NumPy 高级索引来加速切片部分。
然而，排序步骤可能比高级索引更昂贵，并且一般来说，对于我测试的输入，建议的方法往往更快。

import numpy as np


def labeling_unique_argsort(values, labels):
    uniques, counts = np.unique(labels, return_counts=True)
    sorted_values = values[labels.argsort()]
    bound = 0
    result = {}
    for x, c in zip(uniques, counts):
        result[x] = sorted_values[bound:bound + c]
        bound += c
    return result

另一种方法，原则上很简洁(与我提出的方法相同)，但在实践中很慢是使用排序和 itertools.groupby() :

import itertools
from operator import itemgetter


def labeling_groupby(values, labels):
    slicing = labels.argsort()
    sorted_labels = labels[slicing]
    sorted_values = values[slicing]
    del slicing
    result = {}
    for x, g in itertools.groupby(zip(sorted_labels, sorted_values), itemgetter(0)):
        result[x] = np.fromiter(map(itemgetter(1), g), dtype=sorted_values.dtype)
    return result

最后，一种基于 Pandas 的方法，对于较大的输入非常简洁且相当快，但对于较小的输入表现不佳(类似于 @Ehsan's answer):

def labeling_groupby_pd(values, labels):
    df = pd.DataFrame({'values': values, 'labels': labels})
    return df.groupby('labels').values.apply(lambda x: x.values).to_dict()

现在，谈话是廉价的，所以让我们附加一些数字来表示快速和慢速，并为不同的输入大小生成一些图。 K的值上限为 52(英文字母的大小写字母)。当 N远大于 K ，达到上限的概率非常高。
输入是通过以下方式以编程方式生成的:

def gen_input(n, p, labels=string.ascii_letters):
    k = len(labels)
    values = np.arange(n)
    labels = np.array([string.ascii_letters[i] for i in np.random.randint(0, int(k * p), n)])
    return values, labels

并且基准是针对 p 的值生成的来自 (1.0, 0.5, 0.1, 0.05) ，它改变了 K 的最大值.下面的图是指 p值的顺序。

p=1.0 (最多 K = 52)

...并以最快的速度放大


p=0.5 (最多 K = 26)


p=0.1 (最多 K = 5)


p=0.05 (最多 K = 2)

...并以最快的速度放大


可以看到，除了非常小的输入之外，所提出的方法如何优于迄今为止针对测试输入提出的其他方法。
(完整的基准测试可用 here)。

也可以考虑将循环的某些部分移至 Numba/Cython，但我会将其留给感兴趣的读者。