Relax the constraints of floor_sum

atcoder/ac-library#92

Author: kzrnm

Source/AtCoderLibrary/Math/Internal/InternalMath.cs

@@ -152,5 +152,27 @@ public static bool IsPrime(int n)
             }
             return true;
         }
+
+        public static ulong FloorSumUnsigned(ulong n, ulong m, ulong a, ulong b)
+        {
+            ulong ans = 0;
+            while (true)
+            {
+                if (a >= m)
+                {
+                    ans += (n - 1) * n / 2 * (a / m);
+                    a %= m;
+                }
+                if (b >= m)
+                {
+                    ans += n * (b / m);
+                    b %= m;
+                }
+
+                ulong yMax = a * n + b;
+                if (yMax < m) return ans;
+                (n, m, a, b) = (yMax / m, a, m, yMax % m);
+            }
+        }
     }
 }

Source/AtCoderLibrary/Math/MathLib.cs

@@ -572,33 +572,39 @@ public static (long y, long m) CRT(long[] r, long[] m)
 
         #region FloorSum
         /// <summary>
-        /// sum_{i=0}^{<paramref name="n"/>-1} floor(<paramref name="a"/>*i+<paramref name="b"/>/<paramref name="m"/>) を返します。
+        /// sum_{i=0}^{<paramref name="n"/>-1} floor(<paramref name="a"/>*i+<paramref name="b"/>/<paramref name="m"/>) を返します。答えがオーバーフローしたならば  mod2^64 で等しい値を返します。
         /// </summary>
         /// <remarks>
-        /// <para>制約: 0≤<paramref name="n"/>, <paramref name="m"/>≤10^9, 0≤<paramref name="a"/>, <paramref name="b"/>&lt;<paramref name="m"/></para>
-        /// <para>計算量: O(log(n+m+a+b))</para>
+        /// <para>制約:</para>
+        /// <para> 0≤<paramref name="n"/>&lt;2^32</para>
+        /// <para> 1≤<paramref name="m"/>&lt;2^32</para>
+        /// <para>計算量: O(log(m))</para>
         /// </remarks>
         /// <returns></returns>
         public static long FloorSum(long n, long m, long a, long b)
         {
-            long ans = 0;
-            while (true)
+            Contract.Assert(0 <= n && n < (1L << 32));
+            Contract.Assert(1 <= m && m < (1L << 32));
+            var nn = (ulong)n;
+            var mm = (ulong)m;
+            ulong aa, bb;
+            ulong ans = 0;
+            if (a < 0)
             {
-                if (a >= m)
-                {
-                    ans += (n - 1) * n / 2 * (a / m);
-                    a %= m;
-                }
-                if (b >= m)
-                {
-                    ans += n * (b / m);
-                    b %= m;
-                }
-
-                long yMax = a * n + b;
-                if (yMax < m) return ans;
-                (n, m, a, b) = (yMax / m, a, m, yMax % m);
+                var a2 = (ulong)InternalMath.SafeMod(a, m);
+                ans -= nn * (nn - 1) / 2 * ((a2 - (ulong)a) / mm);
+                aa = a2;
+            }
+            else aa = (ulong)a;
+            if (b < 0)
+            {
+                var b2 = (ulong)InternalMath.SafeMod(b, m);
+                ans -= nn * ((b2 - (ulong)b) / mm);
+                bb = b2;
             }
+            else bb = (ulong)b;
+
+            return (long)(ans + InternalMath.FloorSumUnsigned(nn, mm, aa, bb));
         }
         #endregion FloorSum
     }

Test/AtCoderLibrary.Test/Math/FloorSumTest.cs

@@ -15,11 +15,13 @@ public void FloorSum()
             {
                 for (int m = 1; m < 20; m++)
                 {
-                    for (int a = 0; a < 20; a++)
+                    for (int a = -20; a < 20; a++)
                     {
-                        for (int b = 0; b < 20; b++)
+                        for (int b = -20; b < 20; b++)
                         {
-                            MathLib.FloorSum(n, m, a, b).Should().Be(FloorSumNative(n, m, a, b));
+                            var expected = FloorSumNative(n, m, a, b);
+                            MathLib.FloorSum(n, m, a, b).Should()
+                                .Be(expected, "FloorSum({0},{1},{2},{3}) should be {4}", n, m, a, b, expected);
                         }
                     }
                 }

Test/AtCoderLibrary.Test/Utils/MathUtil.cs

@@ -41,12 +41,12 @@ public static long FloorSumNative(long n, long m, long a, long b)
             long sum = 0;
             for (long i = 0; i < n; i++)
             {
-                sum += (a * i + b) / m;
+                long z = a * i + b;
+                sum += (z - InternalMath.SafeMod(z, m)) / m;
             }
             return sum;
         }
 
-
         public static List<int> Factors(int m)
         {
             var result = new List<int>();