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

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

Неравнества с двумя переменными решать вообще сложно. Давайте как-то получим одну переменную. Эталонное медленное решение предполагает, например, перебор всех возможных значений 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;
}

© Habrahabr.ru