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

Есть вот такая, вроде бы, простая задача на литкоде: Дано три числа 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(min(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

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

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

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

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

\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

А поскольку i может принимать все значения до 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 целое. Поскольку x_0 и bочевидно взаимнопросты, то выражение под ceil целое, только для i делящихся на b. Поэтому давайте там везде поставим 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 за логарифм.

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

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

Во-вторых, давайте считать, что k < m. В противном случае можно выделить лишнюю часть, если подставить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}

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

Во-третьих, можно считать, что n \le m. В противном случае, можно разбить сумму на куски, первый несколько раз проходит по 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}

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

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

В правой сумме заменим переменную i = mf+i':

\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через сумму до n' < m.

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

Итак, мы знаем, что GCD(m, k) = 1, k < m, n \le m.

А теперь главная хитрость тут, это применить Hermite’s identity. Я много с ним эксперементировал и самый простой способ получается, если взять заnчисло k, а за xоставшееся \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и t < k. Поэтому значение floor будет или 0 или 1. 1 только в случае \frac{i}{m} + \frac{t}{k} \ge 1. Отсюда можно вывести ограничение на i, при котором что-то ненулевое вообще суммируется: i \ge m - \frac{mt}{k}. Т.к. iцелое, то можно переписать:

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

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

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

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

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

Т.к. tцелое:

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

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

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

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

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

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

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

SumFloor(n, k, m) = (n-m)(k-N')+SumFloor(k, m, k) - SumFloor(N', m, k)

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

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и kвзаимнопросты. Поэтому, после одного рекуррентного перехода мы всегда останемся в том же общем случае, поэтому опять проверять граничные условия не надо.

А база у рекурсии, если k=0, то результат 0. Еще можно остановится, если m=1и выдать сумму арифметической прогресси ik. Или, если n=0, то сумма ничего — ноль.

Подобно алгоритму Эвклида для GCD, эта рекуррентная формула позволяет считать сумму за O(log(max(k,m))). Его еще можно итеративно реализовать и получить O(1)по памяти.

А вот и код:

// Расширенный алгоритм Эвклида
// Возвращает 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;
}

© Habrahabr.ru