cp-book

ecnerwala's competitive programming library

View the Project on GitHub ecnerwala/cp-book

🧪 lattice_cnt.test.cpp

View this file on GitHub · Last update: 2023-11-14 00:10:14-05:00

Depends on

Code

#include <catch2/catch_test_macros.hpp>

#include "lattice_cnt.hpp"

using namespace std;

long long lattice_cnt_slow(long long A, long long B, long long C) {
	using ll = long long;
	ll ans = 0;
	for (ll x = 0; A * x <= C; x++) {
		for (ll y = 0; A * x + B * y <= C; y++) {
			ans++;
		}
	}
	return ans;
}

long long mod_count_range_slow(long long a, long long m, long long clo, long long chi, long long nlo, long long nhi) {
	assert(nlo <= nhi);
	assert(clo <= chi);
	long long ans = 0;
	for (long long i = nlo; i < nhi; i++) {
		for (long long j = clo; j < chi; j++) {
			ans += (((a * i - j) % m) == 0);
		}
	}
	return ans;
}

TEST_CASE("Lattice Count", "[lattice_cnt]") {
	for (int a = 0; a <= 50; a++) {
		for (int b = 0; b <= 10; b++) {
			for (int c = -1; c <= 100; c++) {
				if ((a == 0 || b == 0) && c >= 0) continue;
				INFO("a = " << a);
				INFO("b = " << b);
				INFO("c = " << c);
				REQUIRE(lattice_cnt(a, b, c) == lattice_cnt_slow(a, b, c));
			}
		}
	}
}

TEST_CASE("Mod Count (positive)", "[lattice_cnt]") {
	for (int m = 1; m <= 25; m++) {
		for (int a = 0; a <= m+10; a++) {
			for (int c = 0; c <= m; c++) {
				INFO("a = " << a);
				INFO("m = " << m);
				INFO("c = " << c);
				int trueAns = 0;
				for (int n = 1; n <= m+10; n++) {
					INFO("n = " << n);

					trueAns += (a * (n-1) % m) < c;
					REQUIRE(mod_count(a, m, c, n) == trueAns);
				}
			}
		}
	}
}

TEST_CASE("Mod Count (negatives)", "[lattice_cnt]") {
	for (int m : {1, 2, 3, 5, 8, 13, 21}) {
		for (int a : {-10, 0, 1, 2, 3, 5, m, m+5}) {
			auto cnds = {-37, -2*m-1, -m, -m+1, -m/2, -1, 0, 1, m/2, m+1, 2*m-1, 34};
			INFO("a = " << a);
			INFO("m = " << m);
			for (int clo : cnds) {
				for (int nlo : cnds) {
					INFO("clo = " << clo);
					INFO("nlo = " << nlo);
					REQUIRE(mod_count_range(a, m, clo, 47, nlo, 49) == mod_count_range_slow(a, m, clo, 47, nlo, 49));
				}
			}

			for (int chi : cnds) {
				for (int nhi : cnds) {
					INFO("chi = " << chi);
					INFO("nhi = " << nhi);
					REQUIRE(mod_count_range(a, m, -55, chi, -57, nhi) == mod_count_range_slow(a, m, -55, chi, -57, nhi));
				}
			}
		}
	}
}

#include <catch2/catch_test_macros.hpp>
#include <utility>
#include <cassert>
#line 2 "src/lattice_cnt.test.cpp"

#line 2 "src/lattice_cnt.hpp"

#line 5 "src/lattice_cnt.hpp"

// number of integer solutions to Ax + By <= C and x,y >= 0
inline long long lattice_cnt(long long A, long long B, long long C) {
	using ll = long long;

	assert(A >= 0 && B >= 0);
	if (C < 0) return 0;

	assert(A > 0 && B > 0);
	if (A > B) std::swap(A, B);
	assert(A <= B);

	ll ans = 0;
	while (C >= 0) {
		assert(0 < A && A <= B);

		ll k = B/A;
		ll l = B%A;
		assert(B == k * A + l);

		ll f = C/B;
		ll e = C%B / A;
		ll g = C%B % A;
		assert(C == f * B + e * A + g);
		assert(C == (f * k + e) * A + f * l + g);

		// either x + ky <= f*k+e
		// i.e. 0 <= x <= (f-y) * k + e
		// or x >= fk + e + 1 - ky
		// and Ax + (Ak+l) y <= C = (fk + e + 1) A + fl - A + g
		// Let z = x - (fk + e + 1 - ky)
		// Az + A(fk + e + 1 - ky) + Aky + ly <= C = A (fk + e + 1) + fl - A + g
		// Az + ly <= fl - A + g

		ans += (f+1) * (e+1) + (f+1) * f / 2 * k;

		C = f*l - A + g;
		B = A;
		A = l;
	}
	return ans;
}

// count the number of 0 <= (a * x % m) < c for 0 <= x < n
inline long long mod_count(long long a, long long m, long long c, long long n) {
	assert(m > 0);
	if (n == 0) return 0;

	a %= m; if (a < 0) a += m;

	long long extraC = c / m; c %= m;
	if (c < 0) extraC--, c += m;
	assert(0 <= c && c < m);

	long long ans = extraC * n;

	long long extraN = n / m; n %= m;
	if (n < 0) extraN--, n += m;
	assert(0 <= n && n < m);

	if (extraN) {
		ans += extraN * (lattice_cnt(m, a+m, (a+m) * (m-1)) - lattice_cnt(m, a+m, (a+m) * (m-1) - c));
	}

	if (n) {
		// we want solutions to 0 <= a(N-1-x) - my < c with 0 <= x <= N-1
		// a * (N-1) >= ax + my > a * (N-1) - c
		ans += lattice_cnt(m, a+m, (a+m) * (n-1)) - lattice_cnt(m, a+m, (a+m) * (n-1) - c);
	}

	return ans;
}

inline long long mod_count_range(long long a, long long m, long long clo, long long chi, long long nlo, long long nhi) {
	return mod_count(a, m, chi, nhi) - mod_count(a, m, chi, nlo) - mod_count(a, m, clo, nhi) + mod_count(a, m, clo, nlo);
}
#line 4 "src/lattice_cnt.test.cpp"

using namespace std;

long long lattice_cnt_slow(long long A, long long B, long long C) {
	using ll = long long;
	ll ans = 0;
	for (ll x = 0; A * x <= C; x++) {
		for (ll y = 0; A * x + B * y <= C; y++) {
			ans++;
		}
	}
	return ans;
}

long long mod_count_range_slow(long long a, long long m, long long clo, long long chi, long long nlo, long long nhi) {
	assert(nlo <= nhi);
	assert(clo <= chi);
	long long ans = 0;
	for (long long i = nlo; i < nhi; i++) {
		for (long long j = clo; j < chi; j++) {
			ans += (((a * i - j) % m) == 0);
		}
	}
	return ans;
}

TEST_CASE("Lattice Count", "[lattice_cnt]") {
	for (int a = 0; a <= 50; a++) {
		for (int b = 0; b <= 10; b++) {
			for (int c = -1; c <= 100; c++) {
				if ((a == 0 || b == 0) && c >= 0) continue;
				INFO("a = " << a);
				INFO("b = " << b);
				INFO("c = " << c);
				REQUIRE(lattice_cnt(a, b, c) == lattice_cnt_slow(a, b, c));
			}
		}
	}
}

TEST_CASE("Mod Count (positive)", "[lattice_cnt]") {
	for (int m = 1; m <= 25; m++) {
		for (int a = 0; a <= m+10; a++) {
			for (int c = 0; c <= m; c++) {
				INFO("a = " << a);
				INFO("m = " << m);
				INFO("c = " << c);
				int trueAns = 0;
				for (int n = 1; n <= m+10; n++) {
					INFO("n = " << n);

					trueAns += (a * (n-1) % m) < c;
					REQUIRE(mod_count(a, m, c, n) == trueAns);
				}
			}
		}
	}
}

TEST_CASE("Mod Count (negatives)", "[lattice_cnt]") {
	for (int m : {1, 2, 3, 5, 8, 13, 21}) {
		for (int a : {-10, 0, 1, 2, 3, 5, m, m+5}) {
			auto cnds = {-37, -2*m-1, -m, -m+1, -m/2, -1, 0, 1, m/2, m+1, 2*m-1, 34};
			INFO("a = " << a);
			INFO("m = " << m);
			for (int clo : cnds) {
				for (int nlo : cnds) {
					INFO("clo = " << clo);
					INFO("nlo = " << nlo);
					REQUIRE(mod_count_range(a, m, clo, 47, nlo, 49) == mod_count_range_slow(a, m, clo, 47, nlo, 49));
				}
			}

			for (int chi : cnds) {
				for (int nhi : cnds) {
					INFO("chi = " << chi);
					INFO("nhi = " << nhi);
					REQUIRE(mod_count_range(a, m, -55, chi, -57, nhi) == mod_count_range_slow(a, m, -55, chi, -57, nhi));
				}
			}
		}
	}
}

// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
#include <catch2/catch_test_macros.hpp>
#include <cassert>
// src/lattice_cnt.hpp
inline long long lattice_cnt(long long A,long long B,long long C){
using ll=long long;
assert(A>=0&&B>=0);
if(C<0)return 0;
assert(A>0&&B>0);
if(A>B)std::swap(A,B);
assert(A<=B);
ll ans=0;
while(C>=0){
assert(0<A&&A<=B);
ll k=B/A;
ll l=B%A;
assert(B==k*A+l);
ll f=C/B;
ll e=C%B/A;
ll g=C%B%A;
assert(C==f*B+e*A+g);
assert(C==(f*k+e)*A+f*l+g);
ans+=(f+1)*(e+1)+(f+1)*f/2*k;
C=f*l-A+g;
B=A;
A=l;
}
return ans;
}
inline long long mod_count(long long a,long long m,long long c,long long n){
assert(m>0);
if(n==0)return 0;
a%=m;if(a<0)a+=m;
long long extraC=c/m;c%=m;
if(c<0)extraC--,c+=m;
assert(0<=c&&c<m);
long long ans=extraC*n;
long long extraN=n/m;n%=m;
if(n<0)extraN--,n+=m;
assert(0<=n&&n<m);
if(extraN){
ans+=extraN*(lattice_cnt(m,a+m,(a+m)*(m-1))-lattice_cnt(m,a+m,(a+m)*(m-1)-c));
}
if(n){
ans+=lattice_cnt(m,a+m,(a+m)*(n-1))-lattice_cnt(m,a+m,(a+m)*(n-1)-c);
}
return ans;
}
inline long long mod_count_range(long long a,long long m,long long clo,long long chi,long long nlo,long long nhi){
return mod_count(a,m,chi,nhi)-mod_count(a,m,chi,nlo)-mod_count(a,m,clo,nhi)+mod_count(a,m,clo,nlo);
}
// src/lattice_cnt.test.cpp
using namespace std;
long long lattice_cnt_slow(long long A,long long B,long long C){
using ll=long long;
ll ans=0;
for(ll x=0;A*x<=C;x++){
for(ll y=0;A*x+B*y<=C;y++){
ans++;
}
}
return ans;
}
long long mod_count_range_slow(long long a,long long m,long long clo,long long chi,long long nlo,long long nhi){
assert(nlo<=nhi);
assert(clo<=chi);
long long ans=0;
for(long long i=nlo;i<nhi;i++){
for(long long j=clo;j<chi;j++){
ans+=(((a*i-j)%m)==0);
}
}
return ans;
}
TEST_CASE("Lattice Count","[lattice_cnt]"){
for(int a=0;a<=50;a++){
for(int b=0;b<=10;b++){
for(int c=-1;c<=100;c++){
if((a==0||b==0)&&c>=0)continue;
INFO("a = "<<a);
INFO("b = "<<b);
INFO("c = "<<c);
REQUIRE(lattice_cnt(a,b,c)==lattice_cnt_slow(a,b,c));
}
}
}
}
TEST_CASE("Mod Count (positive)","[lattice_cnt]"){
for(int m=1;m<=25;m++){
for(int a=0;a<=m+10;a++){
for(int c=0;c<=m;c++){
INFO("a = "<<a);
INFO("m = "<<m);
INFO("c = "<<c);
int trueAns=0;
for(int n=1;n<=m+10;n++){
INFO("n = "<<n);
trueAns+=(a*(n-1)%m)<c;
REQUIRE(mod_count(a,m,c,n)==trueAns);
}
}
}
}
}
TEST_CASE("Mod Count (negatives)","[lattice_cnt]"){
for(int m:{1,2,3,5,8,13,21}){
for(int a:{-10,0,1,2,3,5,m,m+5}){
auto cnds={-37,-2*m-1,-m,-m+1,-m/2,-1,0,1,m/2,m+1,2*m-1,34};
INFO("a = "<<a);
INFO("m = "<<m);
for(int clo:cnds){
for(int nlo:cnds){
INFO("clo = "<<clo);
INFO("nlo = "<<nlo);
REQUIRE(mod_count_range(a,m,clo,47,nlo,49)==mod_count_range_slow(a,m,clo,47,nlo,49));
}
}
for(int chi:cnds){
for(int nhi:cnds){
INFO("chi = "<<chi);
INFO("nhi = "<<nhi);
REQUIRE(mod_count_range(a,m,-55,chi,-57,nhi)==mod_count_range_slow(a,m,-55,chi,-57,nhi));
}
}
}
}
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
Back to top page