#853623 (PyPy3) No.922 東北きりきざむたん

提出ソース

結果

問題	No.922 東北きりきざむたん
ユーザー	草苺奶昔
提出日時	2023-03-28 16:24:34
言語	PyPy3 (7.3.15)
結果	WA
実行時間	-
コード長	11,981 bytes
コンパイル時間	202 ms
コンパイル使用メモリ	81,956 KB
実行使用メモリ	202,152 KB
最終ジャッジ日時	2023-10-20 10:33:55
合計ジャッジ時間	24,652 ms
ジャッジサーバーID （参考情報）	judge11 / judge12

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	67 ms 70,192 KB
testcase_01	AC	66 ms 70,192 KB
testcase_02	AC	66 ms 70,192 KB
testcase_03	AC	67 ms 70,192 KB
testcase_04	AC	92 ms 77,392 KB
testcase_05	WA	-
testcase_06	WA	-
testcase_07	AC	109 ms 77,676 KB
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	WA	-
testcase_14	WA	-
testcase_15	AC	380 ms 132,516 KB
testcase_16	WA	-
testcase_17	WA	-
testcase_18	WA	-
testcase_19	WA	-
testcase_20	WA	-
testcase_21	RE	-
testcase_22	RE	-
testcase_23	WA	-
testcase_24	WA	-
testcase_25	AC	997 ms 193,436 KB
testcase_26	AC	1,030 ms 193,064 KB
testcase_27	AC	1,011 ms 192,608 KB
testcase_28	AC	544 ms 155,616 KB
testcase_29	RE	-

権限があれば一括ダウンロードができます

ソースコード

raw source code

from typing import List, Tuple


# class Tree
class LCA_HLD:
    __slots__ = (
        "depth",
        "parent",
        "tree",
        "depthWeighted",
        "lid",
        "rid",
        "_idToNode",
        "_top",
        "_heavySon",
        "_dfn",
    )

    def __init__(self, n: int) -> None:
        self.depth = [0] * n
        self.parent = [-1] * n
        self.tree = [[] for _ in range(n)]
        self.depthWeighted = [0] * n
        self.lid = [0] * n
        self.rid = [0] * n
        self._idToNode = [0] * n
        self._top = [0] * n
        self._heavySon = [0] * n
        self._dfn = 0

    def build(self, root=-1) -> None:
        if root != -1:
            self._build(root, -1, 0, 0)
            self._markTop(root, root)
            return
        n = len(self.tree)
        for i in range(n):
            if self.parent[i] == -1:
                self._build(i, -1, 0, 0)
                self._markTop(i, i)

    def addEdge(self, from_: int, to: int, weight: int) -> None:
        self.tree[from_].append((to, weight))
        self.tree[to].append((from_, weight))

    def addDirectedEdge(self, from_: int, to: int, weight: int) -> None:
        self.tree[from_].append((to, weight))

    def lca(self, u: int, v: int) -> int:
        while True:
            if self.lid[u] > self.lid[v]:
                u, v = v, u
            if self._top[u] == self._top[v]:
                return u
            v = self.parent[self._top[v]]

    def dist(self, u: int, v: int, weighted=True) -> int:
        if weighted:
            return (
                self.depthWeighted[u]
                + self.depthWeighted[v]
                - 2 * self.depthWeighted[self.lca(u, v)]
            )
        return self.depth[u] + self.depth[v] - 2 * self.depth[self.lca(u, v)]

    def kthAncestor(self, root: int, k: int) -> int:
        """k:0-indexed;如果不存在,返回-1.
        kthAncestor(root,0) = root
        """
        if k > self.depth[root]:
            return -1
        while True:
            u = self._top[root]
            if self.lid[root] - k >= self.lid[u]:
                return self._idToNode[self.lid[root] - k]
            k -= self.lid[root] - self.lid[u] + 1
            root = self.parent[u]

    def jump(self, from_: int, to: int, step: int) -> int:
        """
        从 from 节点跳向 to 节点,跳过 step 个节点(0-indexed).
        返回跳到的节点,如果不存在这样的节点,返回-1.
        """
        if step == 1:
            if from_ == to:
                return -1
            if self.isInSubtree(to, from_):
                return self.kthAncestor(to, self.depth[to] - self.depth[from_] - 1)
            return self.parent[from_]
        c = self.lca(from_, to)
        dac = self.depth[from_] - self.depth[c]
        dbc = self.depth[to] - self.depth[c]
        if step > dac + dbc:
            return -1
        if step <= dac:
            return self.kthAncestor(from_, step)
        return self.kthAncestor(to, dac + dbc - step)

    def getPath(self, from_: int, to: int) -> List[int]:
        res = []
        composition = self.getPathDecomposition(from_, to, True)
        for a, b in composition:
            if a <= b:
                res += self._idToNode[a : b + 1]
            else:
                res += self._idToNode[a : b - 1 : -1]
        return res

    def getPathDecomposition(self, from_: int, to: int, vertex: bool) -> List[Tuple[int, int]]:
        """返回沿着`路径顺序`的 [起点,终点] 的 欧拉序 `左闭右闭` 数组.
        注意不一定沿着欧拉序, 但是沿着路径顺序.
        """
        up, down = [], []
        while True:
            if self._top[from_] == self._top[to]:
                break
            if self.lid[from_] < self.lid[to]:
                down.append((self.lid[self._top[to]], self.lid[to]))
                to = self.parent[self._top[to]]
            else:
                up.append((self.lid[from_], self.lid[self._top[from_]]))
                from_ = self.parent[self._top[from_]]
        offset = 1 ^ vertex
        if self.lid[from_] < self.lid[to]:
            down.append((self.lid[from_] + offset, self.lid[to]))
        elif self.lid[to] + offset <= self.lid[from_]:
            up.append((self.lid[from_], self.lid[to] + offset))
        return up + down[::-1]

    def isInSubtree(self, child: int, root: int) -> bool:
        """child是否在root的子树中(child和root不能相等)"""
        return self.lid[root] <= self.lid[child] < self.rid[root]

    def subtreeSize(self, root: int) -> int:
        return self.rid[root] - self.lid[root]

    def id(self, root) -> Tuple[int, int]:
        """返回 root 的欧拉序区间, 左闭右开, 0-indexed."""
        return self.lid[root], self.rid[root]

    def eid(self, u: int, v: int) -> int:
        """返回边 u-v 对应的 欧拉序起点编号, 0-indexed."""
        id1, id2 = self.lid[u], self.lid[v]
        return id1 if id1 > id2 else id2

    def _build(self, cur: int, pre: int, dep: int, dist: int) -> int:
        subSize, heavySize, heavySon = 1, 0, -1
        for next, weight in self.tree[cur]:
            if next != pre:
                nextSize = self._build(next, cur, dep + 1, dist + weight)
                subSize += nextSize
                if nextSize > heavySize:
                    heavySize, heavySon = nextSize, next
        self.depth[cur] = dep
        self.depthWeighted[cur] = dist
        self.parent[cur] = pre
        self._heavySon[cur] = heavySon
        return subSize

    def _markTop(self, cur: int, top: int) -> None:
        self._top[cur] = top
        self.lid[cur] = self._dfn
        self._idToNode[self._dfn] = cur
        self._dfn += 1
        if self._heavySon[cur] != -1:
            self._markTop(self._heavySon[cur], top)
            for next, _ in self.tree[cur]:
                if next != self._heavySon[cur] and next != self.parent[cur]:
                    self._markTop(next, next)
        self.rid[cur] = self._dfn


from collections import defaultdict
from typing import Callable, Dict, Generic, List, TypeVar

T = TypeVar("T")


class RerootingDict(Generic[T]):

    __slots__ = "adjMap", "_allVertex"

    def __init__(self):
        self.adjMap = defaultdict(list)
        self._allVertex = set()

    def addEdge(self, u: int, v: int) -> None:
        self.adjMap[u].append(v)
        self.adjMap[v].append(u)
        self._allVertex |= {u, v}

    def rerooting(
        self,
        e: Callable[[int], T],
        op: Callable[[T, T], T],
        composition: Callable[[T, int, int, int], T],
        root: int,
    ) -> Dict[int, "T"]:
        parents = {v: -1 for v in self._allVertex}
        order = [root]
        stack = [root]
        while stack:
            cur = stack.pop()
            for next in self.adjMap[cur]:
                if next == parents[cur]:
                    continue
                parents[next] = cur
                order.append(next)
                stack.append(next)

        dp1 = {v: e(v) for v in self._allVertex}
        dp2 = {v: e(v) for v in self._allVertex}
        for cur in order[::-1]:
            res = e(cur)
            for next in self.adjMap[cur]:
                if parents[cur] == next:
                    continue
                dp2[next] = res
                res = op(res, composition(dp1[next], cur, next, 0))
            res = e(cur)
            for next in self.adjMap[cur][::-1]:
                if parents[cur] == next:
                    continue
                dp2[next] = op(res, dp2[next])
                res = op(res, composition(dp1[next], cur, next, 0))
            dp1[cur] = res

        for newRoot in order[1:]:
            parent = parents[newRoot]
            dp2[newRoot] = composition(op(dp2[newRoot], dp2[parent]), parent, newRoot, 1)
            dp1[newRoot] = op(dp1[newRoot], dp2[newRoot])
        return dp1



from typing import Dict, Set, Tuple, List

INF = int(4e18)

if __name__ == "__main__":
    from collections import defaultdict
    from typing import DefaultDict, List

    class UnionFindArray:

        __slots__ = ("n", "part", "_parent", "_rank")

        def __init__(self, n: int):
            self.n = n
            self.part = n
            self._parent = list(range(n))
            self._rank = [1] * n

        def find(self, x: int) -> int:
            while x != self._parent[x]:
                self._parent[x] = self._parent[self._parent[x]]
                x = self._parent[x]
            return x

        def union(self, x: int, y: int) -> bool:
            rootX = self.find(x)
            rootY = self.find(y)
            if rootX == rootY:
                return False
            if self._rank[rootX] > self._rank[rootY]:
                rootX, rootY = rootY, rootX
            self._parent[rootX] = rootY
            self._rank[rootY] += self._rank[rootX]
            self.part -= 1
            return True

        def isConnected(self, x: int, y: int) -> bool:
            return self.find(x) == self.find(y)

        def getGroups(self) -> DefaultDict[int, List[int]]:
            groups = defaultdict(list)
            for key in range(self.n):
                root = self.find(key)
                groups[root].append(key)
            return groups

    E = Tuple[int, int]  # (distSum,size) 树中距离之和, 子树大小

    def e(root: int) -> E:
        return (0, 0)

    def op(childRes1: E, childRes2: E) -> E:
        dist1, size1 = childRes1
        dist2, size2 = childRes2
        return (dist1 + dist2, size1 + size2)

    def composition(fromRes: E, parent: int, cur: int, direction: int) -> E:
        """direction: 0: cur -> parent, 1: parent -> cur"""
        dist, size = fromRes
        return (dist + size + 1, size + 1)

    n, m, q = map(int, input().split())
    adjList = [[] for _ in range(n)]
    edges = []
    uf = UnionFindArray(n)
    LCA = LCA_HLD(n)
    for _ in range(m):
        u, v = map(int, input().split())  # 树中保留的边
        u, v = u - 1, v - 1
        adjList[u].append((v, 1))
        adjList[v].append((u, 1))
        edges.append((u, v))
        uf.union(u, v)
        LCA.addEdge(u, v, 1)
    LCA.build(root=-1)

    res = 0
    toCenter = defaultdict(set)  # 每个连通分量内需要到飞机场的点(特殊点)
    for _ in range(q):
        a, b = map(int, input().split())  # 移动的起点和终点
        a, b = a - 1, b - 1
        if uf.isConnected(a, b):
            res += LCA.dist(a, b, weighted=False)
        else:
            toCenter[uf.find(a)].add(a)
            toCenter[uf.find(b)].add(b)

    def getDistSumToSpecials(
        edges: List[Tuple[int, int]], specials: Set[int], root: int
    ) -> Dict[int, int]:
        """求`树`中每个点到特殊点的距离之和."""

        E = Tuple[int, int]  # (distSum,size) 树中距离之和, 子树中特殊点的大小

        def e(root: int) -> E:
            return (0, 0)

        def op(childRes1: E, childRes2: E) -> E:
            dist1, size1 = childRes1
            dist2, size2 = childRes2
            return (dist1 + dist2, size1 + size2)

        def composition(fromRes: E, parent: int, cur: int, direction: int) -> E:
            """direction: 0: cur -> parent, 1: parent -> cur"""
            dist, size = fromRes
            from_ = cur if direction == 0 else parent
            count = 1 if from_ in specials else 0
            return (dist + size + count, size + count)

        R = RerootingDict()
        for u, v in edges:
            R.addEdge(u, v)
        dp = R.rerooting(e, op, composition, root=root)
        return {k: v[0] for k, v in dp.items()}

    groupEdges = defaultdict(list)  # 每个连通分量内的边
    for u, v in edges:
        groupEdges[uf.find(u)].append((u, v))
    groups = uf.getGroups()
    for root, nodes in groups.items():
        dp = getDistSumToSpecials(groupEdges[root], toCenter[root], root)
        res += min(dp.values(), default=0)
    print(res)

yukicoder

結果

テストケース

ソースコード