add lattice_counting

zimpha · zimpha · commit 06f33b111166 · 2018-04-06T14:22:21.000+08:00
diff --git a/mathematics/lattice_counting.cc b/mathematics/lattice_counting.cc
@@ -0,0 +1,15 @@
+using uint64 = unsigned long long;
+using uint128 = __uint128_t;
+
+// \sum_{i=0}^{n-1} \lfloor \frac{a+bi}{m} \rfloor
+uint128 lattice_count(uint64 n, uint64 a, uint64 b, uint64 m) {
+  if (!b) return (uint128)n * (a / m);
+  if (a >= m) return (uint128)n * (a / m) + lattice_count(n, a % m, b, m);
+  if (b >= m) return uint128(n - 1) * n / 2 * (b / m) + lattice_count(n, a, b % m, m);
+  return lattice_count((a + b * n) / m, (a + b * n) % m, m, b);
+}
+
+// \sum_{i=0}^{n-1} i^x (\lfloor \frac{a+bi}{m} \rfloor)^y
+uint128 lattice_count(uint64 n, uint64 a, uint64 b, uint64 m, uint64 x, uint64 y) {
+  if (x == 0 && y == 1) return lattice_count(n, a, b, m);
+}
