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

这样可行，但我认为它不会特别快：

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

获取排列的元素

2 个答案: