Add Pytest and tests for dsu (#5)

.gitignore

@@ -1,3 +1,5 @@
+*.hpp
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -0,0 +1 @@
+作業用ブランチです。一切テストしてないのでバグだらけだと思います。使わないでください。

atcoder/_bit.py

@@ -0,0 +1,15 @@
+def _ceil_pow2(n: int) -> int:
+    x = 0
+    while (1 << x) < n:
+        x += 1
+
+    return x
+
+
+def _bsf(n: int) -> int:
+    x = 0
+    while n % 2 == 0:
+        x += 1
+        n //= 2
+
+    return x

atcoder/_math.py

@@ -0,0 +1,142 @@
+import typing
+
+
+class Barrett:
+    '''
+    Fast moduler by barrett reduction
+    Reference: https://en.wikipedia.org/wiki/Barrett_reduction
+    NOTE: reconsider after Ice Lake
+    '''
+
+    def __init__(self, m: int) -> None:
+        self._m = m
+        self._im = ((1 << 64) - 1) / m + 1
+
+    def umod(self) -> int:
+        return self._m
+
+    def mul(self, a: int, b: int) -> int:
+        '''
+        [1] m = 1
+        a = b = im = 0, so okay
+
+        [2] m >= 2
+        im = ceil(2^64 / m)
+        -> im * m = 2^64 + r (0 <= r < m)
+        let z = a*b = c*m + d (0 <= c, d < m)
+        a*b * im = (c*m + d) * im = c*(im*m) + d*im = c*2^64 + c*r + d*im
+        c*r + d*im < m*m + m*im < m*m + 2^64 + m <= 2^64 + m*(m+1) < 2^64 * 2
+        ((ab * im) >> 64) == c or c + 1
+        '''
+
+        z = a * b
+        x = (z * self._im) >> 64
+        v = z - x * self._m
+        if self._m <= v:
+            v += self._m
+        return v
+
+
+def _is_prime(n: int) -> bool:
+    '''
+    Reference:
+    M. Forisek and J. Jancina,
+    Fast Primality Testing for Integers That Fit into a Machine Word
+    '''
+
+    if n <= 1:
+        return False
+    if n == 2 or n == 7 or n == 61:
+        return True
+    if n % 2 == 0:
+        return False
+
+    d = n - 1
+    while d % 2 == 0:
+        d /= 2
+
+    for a in (2, 7, 61):
+        t = d
+        y = pow(a, t, n)
+        while t != n - 1 and y != 1 and y != n - 1:
+            y = y * y % n
+            t <<= 1
+        if y != n - 1 and t % 2 == 0:
+            return False
+    return True
+
+
+def _inv_gcd(a: int, b: int) -> typing.Tuple[int, int]:
+    a %= b
+    if a == 0:
+        return (b, 0)
+
+    # Contracts:
+    # [1] s - m0 * a = 0 (mod b)
+    # [2] t - m1 * a = 0 (mod b)
+    # [3] s * |m1| + t * |m0| <= b
+    s = b
+    t = a
+    m0 = 0
+    m1 = 1
+
+    while t:
+        u = s // t
+        s -= t * u
+        m0 -= m1 * u  # |m1 * u| <= |m1| * s <= b
+
+        # [3]:
+        # (s - t * u) * |m1| + t * |m0 - m1 * u|
+        # <= s * |m1| - t * u * |m1| + t * (|m0| + |m1| * u)
+        # = s * |m1| + t * |m0| <= b
+
+        s, t = t, s
+        m0, m1 = m1, m0
+
+    # by [3]: |m0| <= b/g
+    # by g != b: |m0| < b/g
+    if m0 < 0:
+        m0 += b // s
+
+    return (s, m0)
+
+
+def _primitive_root(m: int) -> int:
+    if m == 2:
+        return 1
+    if m == 167772161:
+        return 3
+    if m == 469762049:
+        return 3
+    if m == 754974721:
+        return 11
+    if m == 998244353:
+        return 3
+
+    divs = [2] + [0] * 19
+    cnt = 1
+    x = (m - 1) // 2
+    while x % 2 == 0:
+        x //= 2
+
+    i = 3
+    while i * i <= x:
+        if x % i == 0:
+            divs[cnt] = i
+            cnt += 1
+            while x % i == 0:
+                x //= i
+        i += 2
+
+    if x > 1:
+        divs[cnt] = x
+        cnt += 1
+
+    g = 2
+    while True:
+        for i in range(cnt):
+            if pow(g, (m - 1) // divs[i], m) == 1:
+                break
+        else:
+            return g
+        g += 1

atcoder/_queue.py

@@ -0,0 +1,35 @@
+import typing
+
+
+class SimpleQueue:
+    def __init__(self):
+        self._payload = []
+        self._pos = 0
+        self._size = 0
+
+    def reserve(self, n: int) -> None:
+        self._payload += [None] * (n - len(self._payload))
+
+    def __len__(self) -> int:
+        return self._size
+
+    def empty(self) -> bool:
+        return self._size == 0
+
+    def push(self, t: typing.Any) -> None:
+        i = self._pos + self._size
+        if len(self._payload) <= i:
+            self.reserve(i + 1)
+        self._payload[i] = t
+
+    def front(self) -> typing.Any:
+        return self._payload[self._pos]
+
+    def clear(self) -> None:
+        self._payload = []
+        self._pos = 0
+        self._size = 0
+
+    def pop(self) -> None:
+        self._pos += 1
+        self._size -= 1

atcoder/_scc.py

@@ -0,0 +1,98 @@
+import copy
+import typing
+
+
+class CSR:
+    def __init__(
+            self, n: int, edges: typing.List[typing.Tuple[int, int]]) -> None:
+        self.start = [0] * (n + 1)
+        self.elist = [0] * len(edges)
+
+        for e in edges:
+            self.start[e[0] + 1] += 1
+
+        for i in range(1, n + 1):
+            self.start[i] += self.start[i - 1]
+
+        counter = copy.deepcopy(self.start)
+        for e in edges:
+            self.elist[counter[e[0]]] = e[1]
+            counter[e[0]] += 1
+
+
+class SCCGraph:
+    '''
+    Reference:
+    R. Tarjan,
+    Depth-First Search and Linear Graph Algorithms
+    '''
+
+    def __init__(self, n: int) -> None:
+        self._n = n
+        self._edges = []
+
+    def num_vertices(self) -> int:
+        return self._n
+
+    def add_edge(self, from_vertex: int, to_vertex: int) -> None:
+        self._edges.append((from_vertex, to_vertex))
+
+    def scc_ids(self) -> typing.Tuple[int, typing.List[int]]:
+        g = CSR(self._n, self._edges)
+        now_ord = 0
+        group_num = 0
+        visited = []
+        low = [0] * self._n
+        order = [-1] * self._n
+        ids = [0] * self._n
+
+        def dfs(v: int) -> None:
+            nonlocal now_ord
+            nonlocal group_num
+            nonlocal visited
+            nonlocal low
+            nonlocal order
+            nonlocal ids
+
+            low[v] = now_ord
+            order[v] = now_ord
+            now_ord += 1
+            visited.append(v)
+            for i in range(g.start[v], g.start[v + 1]):
+                to = g.elist[i]
+                if order[to] == -1:
+                    dfs(to)
+                    low[v] = min(low[v], low[to])
+                else:
+                    low[v] = min(low[v], order[to])
+
+            if low[v] == order[v]:
+                while True:
+                    u = visited[-1]
+                    visited.pop()
+                    order[u] = self._n
+                    ids[u] = group_num
+                    if u == v:
+                        break
+                group_num += 1
+
+        for i in range(self._n):
+            if order[i] == -1:
+                dfs(i)
+
+        for i in range(self._n):
+            ids[i] = group_num - 1 - ids[i]
+
+        return (group_num, ids)
+
+    def scc(self) -> typing.List[typing.List[int]]:
+        ids = self.scc_ids()
+        group_num = ids[0]
+        counts = [0] * group_num
+        for x in ids[1]:
+            counts[x] += 1
+        groups = [[] for _ in range(group_num)]
+        for i in range(self._n):
+            groups[ids[1][i]].append(i)
+
+        return groups