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

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）：迭代最大需要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())