pow_of_formal_power_seriesを修正

zer0-star · Oct 23, 2023 · 80bbcde · 80bbcde
1 parent cf0c0b4
commit 80bbcde
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/atcoder/extra/math/formal_power_series.nim b/src/atcoder/extra/math/formal_power_series.nim
@@ -477,12 +477,22 @@ proc `{op}`*[T](self: not SparseFormalPowerSeries and not Monomial, r:SparseForm
     mixin pow, init
     var self = self
     deg.revise(self.len)
+    if k == 0:
+      result = initFormalPowerSeries[T](deg)
+      result[0] = T(1)
+      return
     self.resize(deg)
     for i in 0..<deg:
       if not EQUAL(self[i], T(0)):
         let rev = T(1) / self[i]
         result = (((self * rev) shr i).log(deg) * T.init(k)).exp() * (self[i].pow(k))
-        if i * k > deg: return initFormalPowerSeries[T](deg)
+        #if i * k > deg:
+        var p:int
+        if i == 0: p = 0
+        elif k > deg: p = deg + 1
+        else: p = i * k
+        if p > deg:
+          return initFormalPowerSeries[T](deg)
         result = (result shl (i * k)).pre(deg)
         if result.len < deg: result.setlen(deg)
         return