AtCoder Library Python3, PyPy3用

AtCoderから公式スターターパックが遂に待望のリリース！

やべ〜〜〜〜PythonやめてC++に乗り換える方が楽では？

がんばって毎日こつこつ書いて揃えようと思います。名前順とかで書きます。
全部あるかはわかりません。

1. Convolution

mod 998244353に限定したFMT(NTTとも)です。
多項式同士の乗算で各項の係数を求めることができます。nは2冪でa, bの長さより2倍以上大きな値を渡しておき、fmtして得られた数列の同index要素を掛け合わせてifmtすると畳み込みが終了します。重すぎて50万オーダーはさすがに5secでもTLE 7secくらいはいりそう

class FastModuloTransformation:
  def __init__(self, n):
    self.n = n
 
    self.P = 998244353
    self.g = 3
    self.omega = pow(self.g, (self.P - 1) // n, self.P)
  def bitrev(self, a):
    n = self.n
    x = n.bit_length() - 1
    res = a[: ]
    for i in range(n):
      j = 0
      for k in range(x): j |= (i & (1 << k) > 0) << (x - k - 1)
      if i < j: res[i], res[j] = res[j], res[i]
    return res
  def fmt(self, a, base):
    if n == 1: return a
    P = self.P
    omega = base
    A = self.bitrev(a)
    half = pow(omega, n >> 1, P)
    i = n
    m = 1
    while i > 1:
      i >>= 1
      w = pow(omega, i, P)
      wk = 1
      for j in range(m):
        for x in range(j, len(a), 2 * m):
          l = A[x]
          r = wk * A[x + m] % P
          A[x] = (l + r) % P
          A[x + m] = (l + r * half) % P
        wk *= w
        wk %= P
      m <<= 1
    return A
  def ifmt(self, A):
    if self.n == 1: return A
    P = self.P
    nrev = pow(self.n, P - 2, P)
    return [(x * nrev) % P for x in self.fmt(A, pow(self.omega, P - 2, P))]

2. CRT

とど
 

3. Disjoint Set Union

まあただのUnionFindですよね。DSUの方がメジャーな言い方なのかしら。

class DisjointSetUnion():
  def __init__(self, n):
    self.n = n
    self.root = [-1] * (n + 1)
    self.rank = [0] * (n + 1)
  def leader(self, x):
    rt = self.root[x]
    if rt < 0: return x
    else: self.root[x] = self.leader(rt)
    return self.root[x]
  def merge(self, x, y):
    leader = self.leader
    xrt = leader(x)
    yrt = leader(y)
    if xrt == yrt: return
    if self.rank[xrt] > self.rank[yrt]:
      self.root[xrt] += self.root[yrt]
      self.root[yrt] = xrt
    else:
      self.root[yrt] += self.root[xrt]
      self.root[xrt] = yrt
      self.rank[yrt] += self.rank[xrt] == self.rank[yrt]
  def same(self, x, y): return self.leader(x) == self.leader(y)
  def size(self, x):  return -self.root[self.leader(x)]
  def group(self):
    n = self.n
    leader = self.leader
    res = [[] for _ in range(n + 1)]
    for x in range(n + 1): res[leader(x)].append(x)
    return [res[i] for i in range(n + 1) if len(res[i])]

4. Fenwick Tree

indexは[0, n]で使えます。sumは[l, r)の総和を返します。
lowerboundは総和がs以上になるような最小のindexとかを返すやつですがあんまり試していません。

class FenwickTree:
  def __init__(self, n):
    self.n = n
    self.data = [0] * (n + 2)
  def sum(self, l, r):
    s = 0
    while l > 0:
      s -= self.data[l]
      l -= l & -l
    while r > 0:
      s += self.data[r]
      r -= r & -r
    return s
  def add(self, i, x):
    i += 1
    while i <= self.n:
      self.data[i] += x
      i += i & -i
  def lowerbound(self, s):
    x = 0
    y = 0
    for i in range(self.n.bit_length(), -1, -1):
      k = x + (1 << i)
      if k <= self.n and (y + self.data[k] < s):
        y += self.data[k]
        x += 1 << i
    return x + 1

5. Floor Sum

なにこれ？よくわかりませんが一応作りました。

def FloorSum(n, m, a, b):
  # (a * 0 + b) // m + ... + (a * (n - 1) + b) // m
  res = 0
  if a >= m:
    res += (a // m) * n * (n - 1) // 2
    a %= m
  if b >= m:
    res += n * (b // m)
    b %= m
  if a * n + b < m: return res
  u = (a * n + b) // m
  v = (u * m - b)
  res += (n - (v + a - 1) // a) * u + FloorSum(u, a, m, (a - v) % a)
  return res

6. Lazy Segment Tree

とど

7. LCP Array

suffix arrayのLCPが出せます。

def LCPArray(n, s, sa):
  r = [0] * n
  for i in range(n): r[sa[i]] = i
  lcp = [0] * (n - 1)
  x = 0
  for i in range(n):
    if x > 0: x -= 1
    if r[i] == 0: continue
    j = sa[r[i] - 1]
    while max(i, j) + x < n:
      if s[i + x] != s[j + x]: break
      x += 1
    lcp[r[i] - 1] = x
  return lcp 

8. Max Flow

とど

9. Min Cost Flow

とど

10. SCC

とど

11. Segment Tree

とど

12. Sufix Array

induced sortingを使っていません。O(NlogN)くらいです。

def SuffixArray(n, s):
  order = [n - 1 - i for i in range(n)]
  order.sort(key = lambda x: s[x])
  sa = [0] * n
  classes = [0] * n
  for i in range(n):
    sa[i] = order[i]
    classes[i] = S[i]
  k = 1
  while k < n:
    c = classes[: ]
    for i in range(n): 
      if i > 0 and (c[sa[i - 1]] == c[sa[i]]) and ((sa[i - 1] + k) < n) and ((c[sa[i - 1] + k // 2]) == c[sa[i] + k // 2]):
        classes[sa[i]] = classes[sa[i - 1]]
      else: classes[sa[i]] = i
    cc = [i for i in range(n)]
    ss = sa[: ]
    for i in range(n):
      if ss[i] - k >= 0:
        sa[cc[classes[ss[i] - k]]] = ss[i] - k
        cc[classes[ss[i] - k]] += 1
    k <<= 1
  return sa

13. 2-SAT

とど

14. Z-algorithm

ざるごりずむ。各接尾辞に対するLCPがわかる。

def Zalgorithm(S):
  ln = len(S)
  Z = [0] * ln
  Z[0] = ln
  i, j = 1, 0
  while i < ln:
    while i + j < ln and (S[i + j] == S[j]): j += 1
    if j:
      k = 1
      Z[i] = j
      while k < j and (k + Z[k] < j):
        Z[i + k] = Z[k]
        k += 1
      i += k
      j -= k
    else: i += 1
  return Z