Получить элементы перестановки

Question

У меня есть 1D-массив, элементы которого перестановка 0:N, и мне нужно взять первые K элементов этой перестановки

Например, в случае перестановки

0 [[9] 
1  [0] 
2  [1] 
3  [2] 
4  [3] 
5  [4] 
6  [5] 
7  [6] 
8  [7] 
9  [8]] 

первые 3 элементы 9, 8, 7

код

n = start 
    r = zeros (nodeCount, dtype = int) 
    i = 0 
    while (self.nodes[n][direction] != stop): 
     r[i] = n 
     n = self.nodes[n][direction] 
     i+=1 

Мне нужен более быстрый способ, чтобы извлечь элементы из перестановки.

Answer 1

numpy.roll что вы после?

>>> a = np.arange(10) 
>>> b = np.roll(a,1) 
>>> b 
array([9, 0, 1, 2, 3, 4, 5, 6, 7, 8]) 
>>> np.roll(b[::-1],1)[:3] 
array([9, 8, 7]) 

Это последняя строка кода является довольно загадочным, но b[::-1] изменяет массив, np.roll сдвиги его, и [:3] принимает только первые три элемента.

Answer 2

Это работает, но я не думаю, что это будет особенно быстро:

>>> a 
array([9, 0, 1, 2, 3, 4, 5, 6, 7, 8]) 
>>> n = 3 
>>> b = np.empty((n,), dtype=a.dtype) 
>>> b[0] = a[0] 
>>> for k in xrange(1, n): 
...  b[k] = a[b[k-1]] 
... 
>>> b 
array([9, 8, 7])