Внезапно сложная задача на литкоде: Варианты покупки двух товаров / Хабр

Есть вот такая, вроде бы, простая задача на литкоде: Дано три числа 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 - переменные, сколько товаров первого и второго типа мы покупаем. Тогда у нас есть Диофантово неравенство:

$ax+by \le c, x \ge 0, y \ge 0$

Неравнества с двумя переменными решать вообще сложно. Давайте как-то получим одну переменную. Эталонное медленное решение предполагает, например, перебор всех возможных значений x и дальше нахождение количества возможных значений y. Но как это решение ускорять я не придумал. Давайте вместо этого перебирать потраченную сумму:

$ax+by = i, x \ge 0, y \ge 0, 0 \le i \le c$

Это уже Диафантово уравнение с какими-то дополнительными условиями. Но оно весьма стандартное и, поскольку мы уже сократили наибольший общий делитель, можно найти общую формулу для всех его решений:

$x = x_0i+bt \\ y=y_0i-at \\ x_0a+y_0b = 1$

и тут - это какие-то константы, которые можно найти через расширенный алгоритм эвклида. - это свободная переменная, позволяющая прыгать между всеми целочисленными точками на прямой. Сдвиги на вектор $\{b, -a\}$ между соседними целочисленными точками очевидны.

Теперь надо решить дополнительные ограничения на неотрицательность и . Они дают нам ограничения на :

$\frac{-x_0 i}{b} \leq t \leq \frac{y_0 i}{a}$

Поскольку целое, можно расставить округления:

$\lceil \frac{-x_0 i}{b}\rceil \leq t \leq \lfloor \frac{y_0 i}{a} \rfloor$

Чтобы не думать об отрицательных числах, давайте договоримся, что $x_0 \leq 0$ , $y_0 \geq 0$ . Этого можно очевидно добиться. Тут же стоит заметить, что очевидно $\frac{-x_0i}{b} \le \frac{y_0i}{a}$ , поэтому, округляя эти значения вверх и вниз мы в худшем случае получим нижнюю границу на 1 больше верхней границы (если вещественные границы оказались между одними и теми же целыми числами). Поэтому количество целых значений t, удовлетворяющие всем неравнествам (даже если их там таких нет), можно получить так: $\lfloor \frac{y_0 i}{a} \rfloor - \lceil \frac{-x_0 i}{b}\rceil +1$

А поскольку может принимать все значения до c, формула для ответа будет:

$\sum_{i=0}^c \left(\lfloor \frac{y_0 i}{a} \rfloor - \lceil \frac{-x_0 i}{b}\rceil +1\right)$

Тут в перемешку округления вниз и вверх. Давайте выразим ceil через floor. Они обычно отличаются на 1, кроме случаев, когда значение под ceil целое. Поскольку и очевидно взаимнопросты, то выражение под ceil целое, только для делящихся на . Поэтому давайте там везде поставим ceil=floor+1, но запомним, что ровно в $\lfloor\frac{c}{b}\rfloor+1$ итерациях у нас вычитается лишняя единица, их надо будет после суммы назад прибавить. +1 от ceil сократиться с +1 уже под знаком суммы и мы получим:

$\sum_{i=0}^c \left(\lfloor \frac{y_0 i}{a} \rfloor - \lfloor \frac{-x_0 i}{b}\rfloor \right) + \lfloor \frac{c}{b}\rfloor +1$

Теперь разобьем сумму на две и введем вспомогательную функцию:

$SumFloor(n, k, m) = \sum_{i=0}^{n-1} \lfloor \frac{ki}{m}\rfloor$

В итоге, ответ к задаче:

$SumFloor(c+1, y_0, a)-SumFloor(c+1, -x_0, b)+\lfloor\frac{c}{b}\rfloor+1$

Выводим рекуррентную формулу для суммы

Теперь осталась самая малость: научиться считать сумму в SumFloor за логарифм.

Боремся с граничными условиями

Во-первых, можно считать, что и взаимнопросты. Можно их на НОД сократить, вообще не поменяв ни одно слагаемое. Но вообще числа, которые в этой задаче получаются, уже взаимнопросты (см. "очевидно" выше). Но в общем случае один раз в начале подсчитать GCD за логарифм тоже можно - на оценку сложности это не повлияет.

Во-вторых, давайте считать, что . В противном случае можно выделить лишнюю часть, если подставить $k=mf+k', k' = k\%m<m$ в сумму и подсчитать все целое отдельно в арифметической прогрессии:

$\sum_{i=0}^{n-1} \lfloor \frac{(mf+k')i}{m}\rfloor= \sum_{i=0}^{n-1} \left( fi +\lfloor \frac{k'i}{m}\rfloor \right) = \sum_{i=0}^{n-1} \lfloor \frac{k'i}{m}\rfloor + \frac{fn(n-1)}{2}$

Или:

$SumFloor(n, k, m) = SumFloor(n, k\%m, m)+\frac{\lfloor\frac{k}{m}\rfloor n (n-1)}{2}$

Вот мы и выразили сумму с через сумму c .

Во-третьих, можно считать, что $n \le m$ . В противном случае, можно разбить сумму на куски, первый несколько раз проходит по чисел, второй проходит не больше чем m. Пусть $n = mf + n', n' = n\%m < m$ , тогда:

$SumFloor(n, k, m) = \sum_{i=0}^{mf-1} \lfloor \frac{ki}{m}\rfloor + \sum_{i=mf}^{mf+n'-1} \lfloor \frac{ki}{m}\rfloor$

Давайте отдельно рассмтрим левую и правую суммы. В левой сумме воспользуемся равенством $\lfloor a/b \rfloor = \frac{a - a \%b}{b}$ и подсчитаем отдельно не-модули как арифметическую прогрессию:

$\sum_{i=0}^{mf-1} \lfloor \frac{ki}{m}\rfloor = \sum_{i=0}^{mf-1} \frac{ki-(ki)\%m}{m} = \frac{kmf(mf-1)}{2m} - \sum_{i=0}^{mf-1} \frac{(ki)\%m}{m}$

В оставшейся сумме пробегает все возможные остатки по модулю ровно раз. А поскольку и взаимно просты, то очевидно тоже пробегает все возможные остатки по модулю ровнораз, но в каком-то другом порядке. Но порядок нам не важен, можно просто раз просуммировать все остатки от 0 до . Итак, левая сумма сокращается до:

$\frac{kmf(mf-1)}{2m} - f\frac{m(m-1)}{2m} = \frac{f(kmf-k-m+1)}{2}$

В правой сумме заменим переменную :

$\sum_{i=mf}^{mf+n'-1} \lfloor \frac{ki}{m}\rfloor = \sum_{i'=0}^{n'-1} \lfloor \frac{k(mf+i')}{m}\rfloor = \sum_{i'=0}^{n'-1} \lfloor \frac{ki'}{m}\rfloor+kfn'$

Собирая все вместе, помня, что $f = \lfloor \frac{n}{m} \rfloor$ :

$SumFloor(n, k, m) = SumFloor(n\%m, k, m) + \frac{f(kmf-k-m+1)}{2} +kf(n\%m)$

Вот мы и выразили сумму до через сумму до .

Выводим формулу для общего случая

Итак, мы знаем, что , , $n \le m$ .

А теперь главная хитрость тут, это применить Hermite's identity. Я много с ним эксперементировал и самый простой способ получается, если взять зачисло , а за оставшееся $\frac{i}{m}$ :

$\sum_{i=0}^{n-1} \lfloor \frac{ki}{m} \rfloor = \sum_{i=0}^{n-1} \sum_{t=0}^{k-1} \lfloor \frac{i}{m} + \frac{t}{k} \rfloor$

Теперь, как обычно это бывает со вложенными суммами, интегралами или пределами, надо первым делом поменять их порядок:

$SumFloor(n, k, m) = \sum_{t=0}^{k-1} \sum_{i=0}^{n-1} \lfloor \frac{i}{m} + \frac{t}{k} \rfloor$

Следующая догадка состоит в том, чтобы заметить, что под знаком округления вниз стоит сумма из двух чисел, меньших 1 каждое! Ведь $i < n \le m$ и . Поэтому значение floor будет или 0 или 1. 1 только в случае $\frac{i}{m} + \frac{t}{k} \ge 1$ . Отсюда можно вывести ограничение на , при котором что-то ненулевое вообще суммируется: $i \ge m - \frac{mt}{k}$ . Т.к. целое, то можно переписать:

$i \ge m - \lfloor\frac{mt}{k}\rfloor$

Но сумма изначальная по была от 0 до поэтому, чтобы подсчитать только единицы, надо посмотреть, сколько их от $m - \lfloor\frac{mt}{k}\rfloor$ до . Когда эта нижняя граница не превосходит то количество единичных слагаемых будет $n - m + \lfloor \frac{mt}{k} \rfloor$ . Если же эта нижняя граница больше то единиц в сумме вообще не будет, но формула выше даст отрицательное число. Поэтому эту формулу можно использовать, только когда она дает неотрицательное число. Теперь можно вывести условие на , при котором это условие выполняется:

$m - \lfloor \frac{mt}{k} \rfloor \le n$

Можно сгрупировать floor и целые числа по разные стороны. Дальше будет неравнество вида $\lfloor x \rfloor \ge n$ Его можно очевидно заменить на $x \ge n$

$\frac{mt}{k} \ge m-n$ $t \ge \frac{k(m-n)}{m}$

Т.к. целое:

$t \ge \lceil \frac{k(m-n)}{m} \rceil$

Это число, для сокращения выкладок ниже, назовем $N' = \lceil \frac{k(m-n)}{m} \rceil$ .

Итак, пробегает значения от до , а внутри происходит сумма скольки-то единиц. Их количество мы уже считали выше. Итоговая формула:

$SumFloor(n, k, m) = \sum_{t=N'}^{k-1} \left( n - m + \lfloor \frac{mt}{k} \rfloor \right)$

Полезно заметить, что $N' \le k$ (очевидно), поэтому запись выше имеет смысл. Немного спорный вопрос, а что происходит, если , но этот случай нам проблем не составит дальше и в итоге мы получим, что сумма 0 слагаемых равна 0.

Целую часть можно вынести за знак суммы, получив (и даже в случае все работает). Остается почти SumFloor(*, m, k), только вот там сумма не от 0 до какого-то числа, а от N' до какого-то большого значения. Эту сумму можно выразить, как .

Промежуточный итог:

Далее, вспоминая, как мы боролись с граничными случаями, можно заметить, что первая сумма сокращается до:

$SumFloor(k, m, k) = SumFloor(0, m, k) + \frac{(k-1)(m-1)}{2} = \frac{(k-1)(m-1)}{2}$

А вторая до:

$SumFloor(N', m, k) = SumFloor(N', m \% k, k) + \frac{\lfloor\frac{m}{k}\rfloor N' (N'-1)}{2}$

В итоге получаем (напомниаю, что $N' = \lceil \frac{k(m-n)}{m} \rceil$ )

$SumFloor(n, k, m) = \frac{(k-1)(m-1) - \lfloor\frac{m}{k}\rfloor N' (N'-1)}{2} + (m-n)(k-N') - \\ - SumFloor(N', m \% k, k)$

При этом, выше мы уже показали, что $N' \le k$ . Также, очевидно, что $m\%k < k$ . Также, $m\%k$ и взаимнопросты. Поэтому, после одного рекуррентного перехода мы всегда останемся в том же общем случае, поэтому опять проверять граничные условия не надо.

А база у рекурсии, если , то результат 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;
}