Question

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

我想要做的基本上是以X进入子数组X[i]的方式拆分k_ar[i]。执行此操作的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）：迭代最多需要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())

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

1 个答案: