numpy - 获取排列的元素

Question

我有一个一维数组，其元素是 的排列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

这行得通，但我认为它不会特别快：

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

Answer 2

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]只取了前三个元素。