Есть вот такая, вроде бы, простая задача на литкоде: Дано три числа total
- сколько у вас есть денег, cost1
, cost2
- цены двух товаров. Надо подсчитать, сколько всего существует различных способов купить сколько-то этих двух товаров, не выходя из бюджета (значение имеет только общее количество покупок, а не порядок). Иными словами, сколько существет целых неотрицательных пар (x, y), таких что x*cost1+y*cost2 <= total
. Например, имея товары ценами {5, 10} и 20 денег на руках, есть 9 способов потратить деньги: 0, 5, 5+5, 5+5+5, 5+5+5+5, 10, 10+5, 10+5+5, 10+10.
Задача даже помечена как medium, и вообще почти в одну строчку решается, но это если допускать безумно медленное решение за O(total / max(cost1, cost2))
, т.е линейное от входных чисел. А сможете ли вы решить ее сильно быстрее - за O(log(max(cost1, cost2)))
? В этом случае задачка становится вполне себе hard и требует много математики, изобретательности и аккуратности. Если интересно решение, добро пожаловать под кат. Буду рад любым альтернативным решениям. Может кто-то сможет додуматься до похожего решения проще.
Примечание по оформлению
Я буду активно пользоваться "аббревиатурами" в этой статье: всякое нудное объяснение я буду прятать в них. Если вы видите подчеркнутое слово "очевидно", наведите туда мышку и более глубокое объяснение должно всплыть в подсказке.
Сводим задачу к сумме определенного вида
Итак, для экономии букв давайте введем обозначения: a
, b
- цены товаров, c
- сумма на руках. Во-первых, можно считать, что GCD(a,b)=1
, иначе очевидно можно просто сократить все три числа на наибольший общий делитель (в случае с c
, там будет деление нацело). Пусть x
и y
- переменные, сколько товаров первого и второго типа мы покупаем. Тогда у нас есть Диофантово неравенство:
Неравнества с двумя переменными решать вообще сложно. Давайте как-то получим одну переменную. Эталонное медленное решение предполагает, например, перебор всех возможных значений x
и дальше нахождение количества возможных значений y
. Но как это решение ускорять я не придумал. Давайте вместо этого перебирать потраченную сумму:
Это уже Диафантово уравнение с какими-то дополнительными условиями. Но оно весьма стандартное и, поскольку мы уже сократили наибольший общий делитель, можно найти общую формулу для всех его решений:
и
тут - это какие-то константы, которые можно найти через расширенный алгоритм эвклида.
- это свободная переменная, позволяющая прыгать между всеми целочисленными точками на прямой. Сдвиги на вектор
между соседними целочисленными точками очевидны.
Теперь надо решить дополнительные ограничения на неотрицательность и
. Они дают нам ограничения на
:
Поскольку целое, можно расставить округления:
Чтобы не думать об отрицательных числах, давайте договоримся, что ,
. Этого можно очевидно добиться. Тут же стоит заметить, что очевидно
, поэтому, округляя эти значения вверх и вниз мы в худшем случае получим нижнюю границу на 1 больше верхней границы (если вещественные границы оказались между одними и теми же целыми числами). Поэтому количество целых значений t, удовлетворяющие всем неравнествам (даже если их там таких нет), можно получить так:
А поскольку может принимать все значения до
c
, формула для ответа будет:
Тут в перемешку округления вниз и вверх. Давайте выразим ceil через floor. Они обычно отличаются на 1, кроме случаев, когда значение под ceil целое. Поскольку и
очевидно взаимнопросты, то выражение под ceil целое, только для
делящихся на
. Поэтому давайте там везде поставим ceil=floor+1, но запомним, что ровно в
итерациях у нас вычитается лишняя единица, их надо будет после суммы назад прибавить. +1 от ceil сократиться с +1 уже под знаком суммы и мы получим:
Теперь разобьем сумму на две и введем вспомогательную функцию:
В итоге, ответ к задаче:
Выводим рекуррентную формулу для суммы
Теперь осталась самая малость: научиться считать сумму в SumFloor
за логарифм.
Боремся с граничными условиями
Во-первых, можно считать, что и
взаимнопросты. Можно их на НОД сократить, вообще не поменяв ни одно слагаемое. Но вообще числа, которые в этой задаче получаются, уже взаимнопросты (см. "очевидно" выше). Но в общем случае один раз в начале подсчитать GCD за логарифм тоже можно - на оценку сложности это не повлияет.
Во-вторых, давайте считать, что . В противном случае можно выделить лишнюю часть, если подставить
в сумму и подсчитать все целое отдельно в арифметической прогрессии:
Или:
Вот мы и выразили сумму с через сумму c
.
Во-третьих, можно считать, что . В противном случае, можно разбить сумму на куски, первый несколько раз проходит по
чисел, второй проходит не больше чем m. Пусть
, тогда:
Давайте отдельно рассмтрим левую и правую суммы. В левой сумме воспользуемся равенством и подсчитаем отдельно не-модули как арифметическую прогрессию:
В оставшейся сумме пробегает все возможные остатки по модулю
ровно
раз. А поскольку
и
взаимно просты, то очевидно
тоже пробегает все возможные остатки по модулю
ровно
раз, но в каком-то другом порядке. Но порядок нам не важен, можно просто
раз просуммировать все остатки от 0 до
. Итак, левая сумма сокращается до:
В правой сумме заменим переменную :
Собирая все вместе, помня, что :
Вот мы и выразили сумму до через сумму до
.
Выводим формулу для общего случая
Итак, мы знаем, что ,
,
.
А теперь главная хитрость тут, это применить Hermite's identity. Я много с ним эксперементировал и самый простой способ получается, если взять зачисло
, а за
оставшееся
:
Теперь, как обычно это бывает со вложенными суммами, интегралами или пределами, надо первым делом поменять их порядок:
Следующая догадка состоит в том, чтобы заметить, что под знаком округления вниз стоит сумма из двух чисел, меньших 1 каждое! Ведь и
. Поэтому значение floor будет или 0 или 1. 1 только в случае
. Отсюда можно вывести ограничение на
, при котором что-то ненулевое вообще суммируется:
. Т.к.
целое, то можно переписать:
Но сумма изначальная по была от 0 до
поэтому, чтобы подсчитать только единицы, надо посмотреть, сколько их от
до
. Когда эта нижняя граница не превосходит
то количество единичных слагаемых будет
. Если же эта нижняя граница больше
то единиц в сумме вообще не будет, но формула выше даст отрицательное число. Поэтому эту формулу можно использовать, только когда она дает неотрицательное число. Теперь можно вывести условие на
, при котором это условие выполняется:
Можно сгрупировать floor и целые числа по разные стороны. Дальше будет неравнество вида Его можно очевидно заменить на
Т.к. целое:
Это число, для сокращения выкладок ниже, назовем .
Итак, пробегает значения от
до
, а внутри происходит сумма скольки-то единиц. Их количество мы уже считали выше. Итоговая формула:
Полезно заметить, что (очевидно), поэтому запись выше имеет смысл. Немного спорный вопрос, а что происходит, если
, но этот случай нам проблем не составит дальше и в итоге мы получим, что сумма 0 слагаемых равна 0.
Целую часть можно вынести за знак суммы, получив (и даже в случае
все работает). Остается почти SumFloor(*, m, k), только вот там сумма не от 0 до какого-то числа, а от N' до какого-то большого значения. Эту сумму можно выразить, как
.
Промежуточный итог:
Далее, вспоминая, как мы боролись с граничными случаями, можно заметить, что первая сумма сокращается до:
А вторая до:
В итоге получаем (напомниаю, что )
При этом, выше мы уже показали, что . Также, очевидно, что
. Также,
и
взаимнопросты. Поэтому, после одного рекуррентного перехода мы всегда останемся в том же общем случае, поэтому опять проверять граничные условия не надо.
А база у рекурсии, если , то результат 0. Еще можно остановится, если
и выдать сумму арифметической прогресси
. Или, если
, то сумма ничего - ноль.
Подобно алгоритму Эвклида для GCD, эта рекуррентная формула позволяет считать сумму за . Его еще можно итеративно реализовать и получить
по памяти.
А вот и код:
// Расширенный алгоритм Эвклида
// Возвращает GCD(a,b) и переменные x, y, т.ч. a*x+b*y=GCD(a,b)
int GcdEx(int a, int b, int &x, int &y) {
if (b == 0) {
x = 1;
y = 0;
return a;
}
if (b > a) return GcdEx(b, a, y, x);
int xx, yy;
int d = GcdEx(b, a % b, xx, yy);
y = xx - a/b*yy;
x = yy;
return d;
}
// Считает sum i=0.. n-1 floor(i*k/m)
// Ограничения:
// m > 0, 0 <= k < m, 0 <= n <= m, GCD(k, m) = 1
long long FloorSumInternal(long long n, long long k, long long m) {
if (k == 0 || n == 0) return 0;
if (m == 1) return k*n*(n-1)/2;
const long long n1 = ((m-n)*k + m - 1)/m;
long long ans = ((m-1)*(k-1) - (n1-1)*n1*(m/k))/2 + (n-m)*(k-n1);
ans -= FloorSumInternal(n1, m%k, k);
return ans;
}
// Считает sum i=0.. n-1 floor(i*k/m)
// Ограничения только естественные: m > 0, k >= 0, n >= 0
long long FloorSum(long long n, long long k, long long m) {
if (k == 0 || n == 0) return 0;
/*
В общем случае надо сократить GCD.
Но в этой задаче k и m, очевидно, будут взаимнопросты.
const long long d = Gcd(m, k);
m /= d;
k /= d;
*/
if (m == 1) return k*n*(n-1)/2;
if (k >= m) {
return FloorSum(n, k%m, m) + (k/m)*n*(n-1)/2;
}
if (n >= m) {
long long f = n/m;
return FloorSum(n%m, k, m) + f*(k*m*f-k-m+1) / 2 + k*f*(n%m);
}
return FloorSumInternal(n, k, m);
}
// Решение самой задачи.
long long waysToBuyPensPencils(int total, int cost1, int cost2) {
int x0, y0;
int d = GcdEx(cost1, cost2, x0, y0);
cost1 /= d;
cost2 /= d;
total /= d;
if (x0 >= 0) {
int t = x0/cost2 + 1;
x0 -= cost2*t;
y0 += cost1*t;
}
if (y0 <= 0) {
int t = -y0 / cost1 + 1;
x0 -= cost2*t;
y0 += cost1*t;
}
return FloorSum(total+1, y0, cost1) - FloorSum(total+1, -x0, cost2) + total / cost2 + 1;
}