From d9b97893eb5c6e7d5ba1f4a2523eada77e9875c9 Mon Sep 17 00:00:00 2001
From: universato <universato@gmail.com>
Date: Tue, 18 May 2021 16:13:43 +0900
Subject: [PATCH] Relax the constraints of floor_sum

---
 .rubocop.yml                        |  3 ++
 lib/floor_sum.rb                    | 43 +++++++++++++++++++++++------
 lib_lock/ac-library-rb/floor_sum.rb | 43 +++++++++++++++++++++++------
 test/floor_sum_test.rb              |  9 ++++--
 4 files changed, 77 insertions(+), 21 deletions(-)

diff --git a/.rubocop.yml b/.rubocop.yml
index c977c20..e983cfa 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -128,6 +128,9 @@ Style/GlobalVars:
     - 'test/lazy_segtree_test.rb'
 Style/FrozenStringLiteralComment:
   Enabled: false
+Style/InfiniteLoop:
+  Exclude:
+    - 'lib/floor_sum.rb'
 Style/InverseMethods:
   AutoCorrect: false
   Exclude:
diff --git a/lib/floor_sum.rb b/lib/floor_sum.rb
index fc902aa..52e4968 100644
--- a/lib/floor_sum.rb
+++ b/lib/floor_sum.rb
@@ -1,19 +1,44 @@
 def floor_sum(n, m, a, b)
+  raise ArgumentError if n < 0 || m < 1
+
   res = 0
 
-  if a >= m
-    res += (n - 1) * n / 2 * (a / m)
-    a %= m
+  if a < 0
+    a2 = a % m
+    res -= n * (n - 1) / 2 * ((a2 - a) / m)
+    a = a2
   end
 
-  if b >= m
-    res += n * (b / m)
-    b %= m
+  if b < 0
+    b2 = b % m
+    res -= n * ((b2 - b) / m)
+    b = b2
   end
 
-  y_max = a * n + b
-  return res if y_max < m
+  res + floor_sum_unsigned(n, m, a, b)
+end
+
+def floor_sum_unsigned(n, m, a, b)
+  res = 0
+
+  while true
+    if a >= m
+      res += n * (n - 1) / 2 * (a / m)
+      a %= m
+    end
+
+    if b >= m
+      res += n * (b / m)
+      b %= m
+    end
+
+    y_max = a * n + b
+    break if y_max < m
+
+    n = y_max / m
+    b = y_max % m
+    m, a = a, m
+  end
 
-  res += floor_sum(y_max / m, a, m, y_max % m)
   res
 end
diff --git a/lib_lock/ac-library-rb/floor_sum.rb b/lib_lock/ac-library-rb/floor_sum.rb
index 22aa8c7..bf3ee69 100644
--- a/lib_lock/ac-library-rb/floor_sum.rb
+++ b/lib_lock/ac-library-rb/floor_sum.rb
@@ -1,21 +1,46 @@
 module AcLibraryRb
   def floor_sum(n, m, a, b)
+    raise ArgumentError if n < 0 || m < 1
+
     res = 0
 
-    if a >= m
-      res += (n - 1) * n / 2 * (a / m)
-      a %= m
+    if a < 0
+      a2 = a % m
+      res -= n * (n - 1) / 2 * ((a2 - a) / m)
+      a = a2
     end
 
-    if b >= m
-      res += n * (b / m)
-      b %= m
+    if b < 0
+      b2 = b % m
+      res -= n * ((b2 - b) / m)
+      b = b2
     end
 
-    y_max = a * n + b
-    return res if y_max < m
+    res + floor_sum_unsigned(n, m, a, b)
+  end
+
+  def floor_sum_unsigned(n, m, a, b)
+    res = 0
+
+    while true
+      if a >= m
+        res += n * (n - 1) / 2 * (a / m)
+        a %= m
+      end
+
+      if b >= m
+        res += n * (b / m)
+        b %= m
+      end
+
+      y_max = a * n + b
+      break if y_max < m
+
+      n = y_max / m
+      b = y_max % m
+      m, a = a, m
+    end
 
-    res += floor_sum(y_max / m, a, m, y_max % m)
     res
   end
 end
diff --git a/test/floor_sum_test.rb b/test/floor_sum_test.rb
index 5f9f508..6e857b1 100644
--- a/test/floor_sum_test.rb
+++ b/test/floor_sum_test.rb
@@ -7,7 +7,10 @@
 
 def floor_sum_naive(n, m, a, b)
   res = 0
-  n.times { |i| res += (a * i + b) / m }
+  n.times do |i|
+    z = a * i + b
+    res += (z - z % m) / m
+  end
   res
 end
 
@@ -16,8 +19,8 @@ def test_floor_sum
     k = 5
     (0..k).each do |n|
       (1..k).each do |m|
-        (0..k).each do |a|
-          (0..k).each do |b|
+        (-k..k).each do |a|
+          (-k..k).each do |b|
             assert_equal floor_sum_naive(n, m, a, b), floor_sum(n, m, a, b)
           end
         end