按子数组索引列表拆分数组

Question

假设我有一个数组X 和一个索引列表k_ar，其中最大值是K - 1。

我想要做的基本上是以X[i] 进入子数组k_ar[i] 的方式拆分X。 O(n) 的方法如下：

X = [5, 1, 3, 2, 2, 1]

k_ar = [0, 1, 0, 1, 2]

K = max(k_ar) + 1

sub_X = [[] for k in range(K)]

for k, x in zip(k_ar, X):
    sub_X[k].append(x)

虽然这是做这类事情的理想算法，但我想知道 Numpy、Scipy 或任何其他库是否有更快的方法来做这件事。例如，我可以这样做，但它是 O(nK) 而不是 O(n)，因此对于大型 K 来说不是最佳的，尽管在 n 中非常快：

import numpy as np

X = np.ndarray([5, 1, 3, 2, 2, 1], dtype=np.int8)

k_ar = np.ndarray([0, 1, 1, 0, 1, 2], dtype=np.int8)

K = max(k_ar)

sub_X = np.empty(K, dtype=np.ndarray)

for k in range(K):
    sub_X[k] = X[k_ar == k]

那么，有没有一种方法可以在不使用例如的情况下加快速度？ Numba、Cython 还是 PyPy？

Answer 1

您的算法是 O(n)：max 迭代需要 n 步，列表创建迭代需要 n 步，放置迭代也需要 n 步。

另外，我不确定是否有任何理由在迭代期间保持原始列表和索引不变，这意味着您可以通过弹出将内存保持在 n 个元素而不是 2n 个。

最终代码 - O(n) 内存，O(n) CPU：

X = [5, 1, 3, 2, 2, 1]
k_ar = [0, 1, 0, 1, 2]
sub_x = []
while X:
    k = k_ar.pop()
    try:
        sub_x[k].append(X.pop())
    except IndexError:
        sub_x.extend([] for i in range(len(sub_x), k+1))
        sub_x[k].append(X.pop())