CODEWITHSUNDEEP
numpy
alinsoar
alinsoar

Reputation: 15793

Get elements of a permutation

I have an 1D-array whose elements are a permutation of 0:N, and I need to take the first K elements of this permutation

For example, in the case the permutation is

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

The first 3 elements are 9 , 8 , 7

The code is 

    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

I need a faster way to extract the elements from permutation.

Upvotes: 2

Views: 126

Answers (2)

Jaime
Jaime

Reputation: 67437

This works, but I don't think it is going to be especially fast:

>>> 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])

Upvotes: 2

John Vinyard
John Vinyard

Reputation: 13495

Is numpy.roll what you're after?

>>> 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])

That last line of code is pretty cryptic, but b[::-1] reverses the array, np.roll shifts it, and the [:3] only takes the first three elements.

Upvotes: 0

Related Questions