Есть ли способ сократить время, необходимое для нахождения индексов верхних k самых больших элементов из одномерного вектора? - PullRequest
Новая
       7

Есть ли способ сократить время, необходимое для нахождения индексов верхних k самых больших элементов из одномерного вектора?

1 голос
/ 21 июня 2019

Я пытаюсь получить индексы самого большого n и наименьшего n элементов массива в Pytorch, а затем объединить эти элементы в возвращаемый массив, но время, необходимое для этого, крайне необоснованно, так как размер массив увеличивается в геометрической прогрессии.

Я попробовал включенный код, а также отсортировал вектор, чтобы я мог просто взять первый и последний n элементов, но время не уменьшается, и я теряю индексы вектора, что мне и нужно. 

def draw(n, distr):
    return np.concatenate((np.array(distr.topk(k=int(n), largest=True).indices),
                           np.array(distr.topk(k=int(n), largest=False).indices)),
                          axis=0)

Для серии одномерных массивов, различающихся по размеру от длины 10 до длины 2359296; кроме того, переменная n представляет собой целое число, значение которого составляет 1/10 длины массива. Мой компьютер может вычислить индексы самых больших и самых маленьких n элементов всех массивов примерно за 0,5 секунды. Я бы предпочел свести это время к минимуму, насколько это возможно, желательно до менее чем 0,2 секунды

1 Ответ

1 голос
/ 21 июня 2019

Вы должны использовать torch.topk(x,k). 

k=2
x = torch.arange(0,10).resize_((2,5))
print(x)
print("...")
res, ind = torch.topk(x,k)
print(res)

tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]])
...
tensor([[4, 3],
        [9, 8]])
...