Упростите код, чтобы уложилось по времени

Question

ограничение по времени 1 секунда

не проходит большинство тестов , т.к. выполняется дольше(входные данные и за какое время проходит эти тесты не дано)

n, k, s = int(input()), [list(map(lambda x: int(x) - 1, input().split()))[1:] for _ in range(7)], [1, 1, 1, 1, 1, 1, 1]
for _ in range(n - 1):
 s = [sum([s[i] for i in r]) for r in k]
print(sum(s) % 1000000007)

Answer 1

N очень велико. Любой цикл по N не пройдёт по времени. Надо искать другие подходы. Рассмотрим общую задачу. Пусть мы знаем количество песен длины n, начинающихся с разных нот. Как вычислить количество песен длины n + 1? Связь между ними можно выразить как умножение вектора на матрицу v(n + 1) = M * v(n).

Для примера в задаче матрица будет выглядеть так:

    / 0 1 0 0 0 0 0 \  # 1 2
    | 1 0 1 0 0 0 0 |  # 2 3 1
    | 0 1 0 0 0 0 0 |  # 1 2
M = | 0 0 1 0 0 0 0 |  # 1 3
    | 0 0 0 1 1 0 0 |  # 2 4 5
    | 0 0 0 0 0 0 0 |  # 0
    \ 0 0 0 0 0 0 0 /  # 0

Число песен длины один, два и три:

     / 1 \                  / 1 \                  / 2 \
     | 1 |                  | 2 |                  | 2 |
     | 1 |                  | 1 |                  | 2 |
v1 = | 1 |    v2 = M * v1 = | 1 |    v3 = M * v2 = | 1 |
     | 1 |                  | 2 |                  | 3 |
     | 1 |                  | 0 |                  | 0 |
     \ 1 /                  \ 0 /                  \ 0 /

Матричное произведение ассоциативно. Поэтому число песен длины N будет равно vN = M^(N-1) * v1.

Задача сведена к возведению матрицы в степень. Используем быстрое возведение в степень. Это позволит получить степень матрицы за логарифмическое число умножений.

Дополнительно заметим что ответ нам нужем по модулю. Везде в призведениях матриц воспользуемся делением по модулю. Это ограничит длину целых которые умножаются и ещё ускорит решение.

Когда степень матрицы получена её надо умножить на v1 и сложить все числа результата. Получившаяся сумма оказывается равна сумме всех элементов матрицы.

Всё вместе:

def read_line():
    tokens = input().split()[1:]
    notes = [0] * 7
    for i in (int(t) - 1 for t in tokens):
        notes[i] = 1
    return tuple(notes)

def mul(a, b):
    return tuple(
        tuple(
            sum(a[i][k] * b[k][j] for k in range(7)) % 1000000007
            for j in range(7)
        )
        for i in range(7)
    )

unit = tuple(
    tuple(
        1 if i == j else 0
        for j in range(7)
    )
    for i in range(7)
)

def pow(a, n):
    if n == 0:
        return unit
    if n == 1:
        return a
    if n % 2 == 0:
        p2 = pow(a, n // 2)
        return mul(p2, p2)
    return mul(a, pow(a, n - 1))

n = int(input())
m = tuple(read_line() for _ in range(7))
p = pow(m, n - 1)
print(sum(map(sum, p)) % 1000000007)

БЛОГ НА HUSL

Упростите код, чтобы уложилось по времени

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