Как можно оптимизировать алгоритм на питоне?

Question

Решаю задачи прошлых лет от яндекса и всё никак не могу пройти тест с большим набором данных, дропается по времени. Сама задача

Мой код:

n = int(input())
orders = []
for i in range(n):
    start = list(map(int, input().split()))
    orders.append(start)

q = int(input())
queries = []
for i in range(q):
    start = list(map(int, input().split()))
    queries.append(start)


results = []
for query in queries:
    query_type = query[-1]
    if query_type == 1:
        start_time = query[0]
        end_time = query[1]
        total_cost = 0
        for order in orders:
            order_start_time = order[0]
            if order_start_time >= start_time and order_start_time <= end_time:
                total_cost += order[2]
        results.append(total_cost)
    elif query_type == 2:
        start_time = query[0]
        end_time = query[1]
        total_duration = 0
        for order in orders:
            order_end_time = order[1]
            if order_end_time >= start_time and order_end_time <= end_time:
                total_duration += order[1] - order[0]
        results.append(total_duration)

print(*results)

Пример входных и выходных данных

Ввод

1
10 100 1000
6
1 10 1
1 10 2
10 100 1
10 100 2
100 1000 1
100 1000 2

Вывод

1000 0 1000 90 0 90

Ввод

Вывод

10 12 3

Answer 1

n = int(input())
orders = [list(map(int, input().split())) for _ in range(n)]

q = int(input())
queries = [list(map(int, input().split())) for _ in range(q)]

results = []
for query in queries:
    query_type = query[-1]
    start_time = query[0]
    end_time = query[1]
    total = 0

    if query_type == 1:
        for order in orders:
            if start_time <= order[0] <= end_time:
                total += order[2]
    elif query_type == 2:
        for order in orders:
            if start_time <= order[1] <= end_time:
                total += order[1] - order[0]

    results.append(total)

print(*results)

Заменены циклы на list comprehensions для создания списка orders и queries, так как это более эффективно и читаемо. Также объединение условия для подсчета total в один цикл, чтобы избежать дублирования кода и ускорить выполнение программы.

Answer 2

Заметание. События упорядочены по времени и типу события.

тип события	значение
0	начало периода запроса
1	начало или конец заказа
2	конец периода запроса

Типы упорядочены так, чтобы период запроса оказался замкнутым интервалом.

Статус: накопленная стоимость totals[0] и накопленная длительность totals[1].

событие	действие
начало заказа	Увеличение накопленной стоимости.
конец заказа	Увеличение накопленной длительности.
начало запроса типа 1	Сохранение накопленной стоимости в заказе со знаком минус.
конец запроса типа 1	Добавление накопленной стоимости в заказ. В заказе оказывается разность накопленных стоимостей в начале и конце заказа.
начало запроса типа 2	Cохранение накопленной длительности в заказе со знаком минус.
конец запроса типа 2	Добавление накопленной длительности в заказ. В заказе оказывается разность накопленных длительностей в начале и конце заказа.

Сложность решения O((n + q)log(n + q)). Львиная доля времени уходит на сортировку. Само заметание линейное. Пространственная сложность O(n + q).

events = []
totals = [0, 0]


def make_total_increment(i, v):

    def f():
        totals[i] += v

    return f


n = int(input())

for _ in range(n):
    s, e, c = map(int, input().split())
    events.append((s, 1, make_total_increment(0, c)))
    events.append((e, 1, make_total_increment(1, e - s)))

q = int(input())

queries = [0] * q


def make_query_increment(i, j):

    def f():
        queries[i] += totals[j]

    return f


def make_query_decrement(i, j):

    def f():
        queries[i] -= totals[j]

    return f


for i in range(q):
    s, e, t = map(int, input().split())
    if t == 1:
        events.append((s, 0, make_query_decrement(i, 0)))
        events.append((e, 2, make_query_increment(i, 0)))
    else:
        events.append((s, 0, make_query_decrement(i, 1)))
        events.append((e, 2, make_query_increment(i, 1)))

for _, _, f in sorted(events, key=lambda v: v[:2]):
    f()

print(*queries)

Answer 3

Так как запросы касаются сумм на интервалах, упорядочим исходные значения по времени и сохраним накопительные суммы. Для вычисления суммы на интервале нужно отыскать левый и правый концы интервала и вычесть накопительные суммы на концах.

Сложность алгоритма O((n + q)logn). nlogn - сортировка заказов, qlogn - исполнение запросов. Пространственная сложность O(n).

import bisect


def make_range_sum(points):
    x, y = zip(*sorted(points))
    
    ys = [0]
    for yi in y:
        ys.append(ys[-1] + yi)

    def range_sum(a, b):
        i = bisect.bisect_left(x, a)
        j = bisect.bisect_right(x, b)
        return ys[j] - ys[i]

    return range_sum


costs = []
durations = []

n = int(input())

for _ in range(n):
    s, e, c = map(int, input().split())
    costs.append((s, c))
    durations.append((e, e - s))

range_sums = None, make_range_sum(costs), make_range_sum(durations)

for i in range(int(input())):
    if i > 0:
        print(end=' ')
    s, e, t = map(int, input().split())
    print(range_sums[t](s, e), end='')
print()

БЛОГ НА HUSL

Как можно оптимизировать алгоритм на питоне?

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