Помогите избавиться от Time Limit Execeeded

Подскажите, пожалуйста по задаче о порядковых статистиках. Код работает, ответ получается правильный. Но каждый раз система автопроверки выдает TL. Что нужно исправить?

Требуется уметь быстро находить k-ю порядковую статистику последовательности — то есть элемент, который после сортировки этой последовательности по неубыванию будет стоять в ней на k-м месте.

Есть две замкнутые в кольцо ленты, на каждой из которых записана последовательность чисел: числа a_1, a_2, . . ., a_n на первой ленте и числа b_1, b_2, . . ., b_n на второй. Назовём p-м циклическим cдвигом положение, в котором a_1 находится над b_(p+1), a_2 над b_(p+2), . . ., a_(n−p) над bn, a_(n−p+1) над b_1, . . ., a_(n−1) над b_(p−1) и a_n над b_p. Рассмотрим последовательность дробей:

(a_1)/(b_(p+1)), (a_2)/(b_(p+2)), . . . , (a_(n−p))/b_n, (a_(n−p+1))/b_1, . . . ,(a_(n−1))/b_((p−1)), (a_n)/(b_p)

и обозначим k-ю порядковую статистику этой последовательности как (s_k)^p.

По данным числам на лентах, а также числу k, найдите последовательность (s_k)^p, (s_k)^2, . . ., (s_k)^n.

Формат входных данных:

В первой строке входных данных заданы через пробел два целых числа n и k — размер лент и номер порядковой статистики (1 ≤ k ≤ n ≤ 5000). Во второй строке записаны n чисел a_1, a_2, . . ., a_n — содержимое ленты с числителями (1 ≤ a_i ≤ 10^9). В третьей строке записаны n чисел b_1, b_2, . . ., b_n — содержимое ленты со знаменателями (1 ≤ b_i ≤ 10^9).

Формат выходных данных:

В первой строке выведите n чисел через пробел — k-е порядковые статистики (s_k)^1, (s_k)^2, . . ., (s_k)^n. Каждое число должно быть представлено в виде несократимой дроби

Примеры

стандартный ввод стандартный вывод
3 1
1 2 3
1 2 3
1/2 1/3 1/1
2 2
4 9
4 2
9/4 9/2
3 2
2 2 2
3 4 5
1/2 1/2 1/2
from fractions import Fraction
import heapq

def kth_statistics(n, k, a, b):
    results = []

    for p in range(1, n + 1):
        fractions = []
        for i in range(n):
            numerator = a[i]
            if i + p < n:
                denominator = b[i + p]
            else:
                denominator = b[(i + p) % n]

            fraction = Fraction(numerator, denominator)
            fractions.append(fraction)

        
        kth_statistic = heapq.nsmallest(k, fractions)[-1]
        results.append(kth_statistic)

    return results

n, k = map(int, input().strip().split())
a = list(map(int, input().strip().split()))
b = list(map(int, input().strip().split()))

result = kth_statistics(n, k, a, b)

print(' '.join(f"{frac.numerator}/{frac.denominator}" for frac in result))

Ответы (1 шт):

Автор решения: MBo

Несколько ускорить можно, используя кортежи для дробей (gord1402 уже предложил подобное в комментариях), и алгоритм QuickSelect для выбора нужной статистики (это побыстрее для больших k, т.к. подход с кучей генерирует длинный список).

def kth_statistics(n, k, a, b):
    results = []

    for p in range(1, n + 1):
        fractions = [(a[i], b[(i + p) % n]) for i in range(n)]

        kth_statistic = QuickSelect(fractions, k, 0, n-1)
        results.append(kth_statistic)

    return results

Сравнение в QuickSelect делается так:

if A[j][0]*pivot[1] < pivot[0]* A[j][1]:

Получается примерно на порядок быстрее, однако алгоритм все равно остается (в среднем) квадратичным, для n=5000 у меня работает на случайных наборах данных до 10 секунд (исходный - от 30 до 230 секунд в зависимости от k), так что, вероятно, требуется найти математическую изюминку для радикального изменения алгоритма.

→ Ссылка