test/math/floor-sum/yosupo-Sum-of-Floor-of-Linear.test.cpp

• View this file on GitHub
• Last update: 2023-12-03 13:41:56+08:00
• Problem: https://judge.yosupo.jp/problem/sum_of_floor_of_linear
• Link: View error logs on GitHub Actions

Depends on

• Sum of floor linear ($\sum_{i = 0}^{n - 1} \lfloor \frac{ai + b}{m} \rfloor$) (library/math/floor-sum.hpp)
• library/math/safe-mod.hpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/sum_of_floor_of_linear"

#include <iostream>

#include "../../../library/math/floor-sum.hpp"

using namespace std;
using namespace felix;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	int tt;
	cin >> tt;
	while(tt--) {
		int n, m, a, b;
		cin >> n >> c >> a >> b;
		cout << floor_sum(n, a, b, c) << "\n";
	}
	return 0;
}

#line 1 "test/math/floor-sum/yosupo-Sum-of-Floor-of-Linear.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/sum_of_floor_of_linear"

#include <iostream>

#line 2 "library/math/floor-sum.hpp"
#include <algorithm>

#include <cassert>

#line 2 "library/math/safe-mod.hpp"

namespace felix {

namespace internal {

template<class T>
constexpr T safe_mod(T x, T m) {
	x %= m;
	if(x < 0) {
		x += m;
	}
	return x;
}

} // namespace internal


} // namespace felix
#line 5 "library/math/floor-sum.hpp"

namespace felix {

// sum_{i = 0}^{n - 1} floor((ai + b) / c) in O(a + b + c + n)

long long floor_sum(long long n, long long a, long long b, long long c) {
	assert(0 <= n && n < (1LL << 32));
	assert(1 <= c && c < (1LL << 32));
	unsigned long long ans = 0;
	if(a < 0) {
		unsigned long long a2 = internal::safe_mod(a, c);
		ans -= 1ULL * n * (n - 1) / 2 * ((a2 - a) / c);
		a = a2;
	}
	if(b < 0) {
		unsigned long long b2 = internal::safe_mod(b, c);
		ans -= 1ULL * n * ((b2 - b) / c);
		b = b2;
	}
	unsigned long long N = n, C = c, A = a, B = b;
	while(true) {
		if(A >= C) {
			ans += N * (N - 1) / 2 * (A / C);
			A %= C;
		}
		if(B >= C) {
			ans += N * (B / C);
			B %= C;
		}
		unsigned long long y_max = A * N + B;
		if(y_max < C) {
			break;
		}
		N = y_max / C;
		B = y_max % C;
		std::swap(C, A);
	}
	return ans;
}

} // namespace felix

#line 5 "test/math/floor-sum/yosupo-Sum-of-Floor-of-Linear.test.cpp"
using namespace std;
using namespace felix;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	int tt;
	cin >> tt;
	while(tt--) {
		int n, m, a, b;
		cin >> n >> c >> a >> b;
		cout << floor_sum(n, a, b, c) << "\n";
	}
	return 0;
}

Back to top page