lattice_cnt.hpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <utility> | ||
| 4 | #include <cassert> | ||
| 5 | |||
| 6 | // number of integer solutions to Ax + By <= C and x,y >= 0 | ||
| 7 | 994569 | inline long long lattice_cnt(long long A, long long B, long long C) { | |
| 8 | 994569 | using ll = long long; | |
| 9 | |||
| 10 | 994569 | assert(A >= 0 && B >= 0); | |
| 11 |
2/4✓ Branch 5 → 6 taken 959256 times.
✓ Branch 5 → 30 taken 35313 times.
✗ Branch 6 → 7 not taken.
✗ Branch 6 → 8 not taken.
|
994569 | if (C < 0) return 0; |
| 12 | |||
| 13 | 959256 | assert(A > 0 && B > 0); | |
| 14 |
2/4✓ Branch 9 → 10 taken 44945 times.
✓ Branch 9 → 11 taken 914311 times.
✗ Branch 12 → 13 not taken.
✗ Branch 12 → 14 not taken.
|
959256 | if (A > B) std::swap(A, B); |
| 15 | 959256 | assert(A <= B); | |
| 16 | |||
| 17 | ✗ | ll ans = 0; | |
| 18 |
2/4✗ Branch 28 → 17 not taken.
✗ Branch 28 → 29 not taken.
✓ Branch 29 → 13 taken 1974079 times.
✓ Branch 29 → 30 taken 959256 times.
|
2933335 | while (C >= 0) { |
| 19 | 1974079 | assert(0 < A && A <= B); | |
| 20 | |||
| 21 |
1/2✗ Branch 16 → 17 not taken.
✓ Branch 16 → 18 taken 1974079 times.
|
1974079 | ll k = B/A; |
| 22 | 1974079 | ll l = B%A; | |
| 23 | 1974079 | assert(B == k * A + l); | |
| 24 | |||
| 25 |
1/2✗ Branch 20 → 21 not taken.
✓ Branch 20 → 22 taken 1974079 times.
|
1974079 | ll f = C/B; |
| 26 |
1/2✗ Branch 22 → 23 not taken.
✓ Branch 22 → 24 taken 1974079 times.
|
1974079 | ll e = C%B / A; |
| 27 | 1974079 | ll g = C%B % A; | |
| 28 | 1974079 | assert(C == f * B + e * A + g); | |
| 29 | 1974079 | assert(C == (f * k + e) * A + f * l + g); | |
| 30 | |||
| 31 | // either x + ky <= f*k+e | ||
| 32 | // i.e. 0 <= x <= (f-y) * k + e | ||
| 33 | // or x >= fk + e + 1 - ky | ||
| 34 | // and Ax + (Ak+l) y <= C = (fk + e + 1) A + fl - A + g | ||
| 35 | // Let z = x - (fk + e + 1 - ky) | ||
| 36 | // Az + A(fk + e + 1 - ky) + Aky + ly <= C = A (fk + e + 1) + fl - A + g | ||
| 37 | // Az + ly <= fl - A + g | ||
| 38 | |||
| 39 | 1974079 | ans += (f+1) * (e+1) + (f+1) * f / 2 * k; | |
| 40 | |||
| 41 | 1974079 | C = f*l - A + g; | |
| 42 | 1974079 | B = A; | |
| 43 | 1974079 | A = l; | |
| 44 | } | ||
| 45 | ✗ | return ans; | |
| 46 | } | ||
| 47 | |||
| 48 | // count the number of 0 <= (a * x % m) < c for 0 <= x < n | ||
| 49 | 337012 | inline long long mod_count(long long a, long long m, long long c, long long n) { | |
| 50 | 337012 | assert(m > 0); | |
| 51 |
2/3✓ Branch 4 → 5 taken 333172 times.
✗ Branch 4 → 6 not taken.
✓ Branch 4 → 29 taken 3840 times.
|
337012 | if (n == 0) return 0; |
| 52 | |||
| 53 |
3/6✗ Branch 5 → 6 not taken.
✓ Branch 5 → 7 taken 333172 times.
✗ Branch 6 → 7 not taken.
✗ Branch 6 → 8 not taken.
✓ Branch 7 → 8 taken 4416 times.
✓ Branch 7 → 9 taken 328756 times.
|
333172 | a %= m; if (a < 0) a += m; |
| 54 | |||
| 55 |
1/2✗ Branch 9 → 10 not taken.
✓ Branch 9 → 11 taken 333172 times.
|
333172 | long long extraC = c / m; c %= m; |
| 56 |
2/4✗ Branch 8 → 9 not taken.
✗ Branch 8 → 10 not taken.
✓ Branch 11 → 12 taken 22080 times.
✓ Branch 11 → 13 taken 311092 times.
|
333172 | if (c < 0) extraC--, c += m; |
| 57 | 333172 | assert(0 <= c && c < m); | |
| 58 | |||
| 59 | 333172 | long long ans = extraC * n; | |
| 60 | |||
| 61 |
1/2✗ Branch 15 → 16 not taken.
✓ Branch 15 → 17 taken 333172 times.
|
333172 | long long extraN = n / m; n %= m; |
| 62 |
2/4✗ Branch 14 → 15 not taken.
✗ Branch 14 → 16 not taken.
✓ Branch 17 → 18 taken 23040 times.
✓ Branch 17 → 19 taken 310132 times.
|
333172 | if (n < 0) extraN--, n += m; |
| 63 | 333172 | assert(0 <= n && n < m); | |
| 64 | |||
| 65 |
2/4✗ Branch 20 → 21 not taken.
✗ Branch 20 → 24 not taken.
✓ Branch 21 → 22 taken 162764 times.
✓ Branch 21 → 25 taken 170408 times.
|
333172 | if (extraN) { |
| 66 | 162764 | ans += extraN * (lattice_cnt(m, a+m, (a+m) * (m-1)) - lattice_cnt(m, a+m, (a+m) * (m-1) - c)); | |
| 67 | } | ||
| 68 | |||
| 69 |
2/4✗ Branch 24 → 25 not taken.
✗ Branch 24 → 28 not taken.
✓ Branch 25 → 26 taken 308990 times.
✓ Branch 25 → 29 taken 24182 times.
|
333172 | if (n) { |
| 70 | // we want solutions to 0 <= a(N-1-x) - my < c with 0 <= x <= N-1 | ||
| 71 | // a * (N-1) >= ax + my > a * (N-1) - c | ||
| 72 | 308990 | ans += lattice_cnt(m, a+m, (a+m) * (n-1)) - lattice_cnt(m, a+m, (a+m) * (n-1) - c); | |
| 73 | } | ||
| 74 | |||
| 75 | ✗ | return ans; | |
| 76 | } | ||
| 77 | |||
| 78 | 16128 | inline long long mod_count_range(long long a, long long m, long long clo, long long chi, long long nlo, long long nhi) { | |
| 79 | 16128 | 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); | |
| 80 | } | ||
| 81 |