Relax the constraints of floor_sum

.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:

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

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

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