Решение задач с помощью матриц

В этом посте рассмотрю задачу с таким условием:

У вас есть массив A длины K, заполненный числами. Каждый следующий элемент массива является суммой K предыдущих элементов. Найти значение элемента, который будет на позиции N. Так как ответ может быть слишком большой, его надо вывести по модулю 1e9+7.

Решение №1

Данное решение является самым примитивным и заключается в явном получении массива A длины N и выводе последнего элемента.

На каждом шаге нужно просуммировать K предыдущих элементов, а таких шагов будет N-K.

Итоговая асимптотика — O(K*(N-K))

Решение №2

Заметим, что используя префиксные суммы можно не суммировать каждый раз K предыдущих элементов, а получать это значение за O(1).

Итоговая асимптотика — O(N)

Решение №3

Теперь используем матрицы.

Чтобы быстро посчитать N-ый элемент массива A, нужно найти такую матрицу, что:

(A[i-k+1], A[i-k+2], ... , A[i]) * X = (A[i-k+2], A[i-k+3], ... , A[i+1])

Матрица X строится следующим методом:

Матрица будет иметь размеры K*K и изначально состоять из нулей, чтобы после умножения матрицы размера K*1 получилась матрица размера K*1
Посмотрим на квадрат размера (K-1)*(K-1) с левым нижним углом в левом нижнем углу матрицы K*K. Заменим все элементы главной диагонали на 1.
Все элементы последнего столбца матрицы K*K заменим на 1.

Докажу корректность данного построения:

После умножения матрицу X получится матрица:

(A[i-k+2], A[i-k+3], ... , A[i-k+1] + A[i-k+2] + ... + A[i])

Так как по условию следующий элемент массива является суммой K предыдущих элементов, то последний элемент как раз равен A[i+1]. Т.е. была получена нужная матрица.

Использование

Необходимо создать начальную матрицу 1*K элементов, далее возвести матрицу X в степень (N-K) с помощью бинарного возведения в степень по модулю 1e9+7.

В конце надо умножить начальную матрицу на полученную матрицу X^(N-K) по модулю 1e9+7 и получить ответ. Он будет находится в K-ом элементе итоговой матрицы.

Итоговая асимптотика — O(logN * (K^3))

Избавиться от куба в степени можно с помощью более быстрых алгоритмов перемножения матриц.

Примеры матриц для разных K:

K = 1:

(1)

K = 2:

(0 1)

(1 1)

K = 3:

(0 0 1)

(1 0 1)

(0 1 1)

K = 4:

(0 0 0 1)

(1 0 0 1)

(0 1 0 1)

(0 0 1 1)

K = 5:

(0 0 0 0 1)

(1 0 0 0 1)

(0 1 0 0 1)

(0 0 1 0 1)

(0 0 0 1 1)

P.S. при K = 2 данная задача сильно похожа на идею чисел Фибоначчи

Итоги

Выбор варианта решения будет исходить из ограничений в задаче, при больших K и небольших N следует использовать 2-ой вариант, при обратном случае, когда большое N и небольшое K следует использовать матрицы.