#833207 (PyPy3) No.529 帰省ラッシュ

提出ソース

結果

問題	No.529 帰省ラッシュ
ユーザー	terasa
提出日時	2023-01-14 12:47:28
言語	PyPy3 (7.3.15)
結果	AC
実行時間	2,133 ms / 4,500 ms
コード長	11,186 bytes
コンパイル時間	529 ms
コンパイル使用メモリ	82,200 KB
実行使用メモリ	201,664 KB
最終ジャッジ日時	2024-06-07 14:40:24
合計ジャッジ時間	22,887 ms
ジャッジサーバーID （参考情報）	judge1 / judge3

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	69 ms 71,160 KB
testcase_01	AC	68 ms 70,460 KB
testcase_02	AC	69 ms 70,892 KB
testcase_03	AC	69 ms 71,804 KB
testcase_04	AC	226 ms 82,268 KB
testcase_05	AC	228 ms 82,084 KB
testcase_06	AC	242 ms 82,520 KB
testcase_07	AC	238 ms 82,680 KB
testcase_08	AC	1,732 ms 187,612 KB
testcase_09	AC	1,588 ms 179,488 KB
testcase_10	AC	1,714 ms 181,968 KB
testcase_11	AC	1,784 ms 184,224 KB
testcase_12	AC	1,332 ms 156,876 KB
testcase_13	AC	1,067 ms 201,664 KB
testcase_14	AC	1,189 ms 192,488 KB
testcase_15	AC	2,133 ms 190,368 KB
testcase_16	AC	2,119 ms 190,808 KB
testcase_17	AC	1,792 ms 200,764 KB
testcase_18	AC	1,666 ms 199,660 KB
testcase_19	AC	1,654 ms 197,276 KB

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

ソースコード

raw source code

from typing import List, Tuple, Callable, TypeVar, Optional
import sys
import itertools
import heapq
import bisect
import math
from collections import deque, defaultdict
from functools import lru_cache, cmp_to_key

input = sys.stdin.readline

if __file__ != 'prog.py':
    sys.setrecursionlimit(10 ** 6)


def readints(): return map(int, input().split())
def readlist(): return list(readints())
def readstr(): return input()[:-1]
def readlist1(): return list(map(lambda x: int(x) - 1, input().split()))


class SegTree:
    def __init__(self, N, func, e):
        self.N = N
        self.func = func
        self.X = [e] * (N << 1)
        self.e = e

    def build(self, seq):
        for i in range(self.N):
            self.X[self.N + i] = seq[i]
        for i in range(self.N)[::-1]:
            self.X[i] = self.func(self.X[i << 1], self.X[i << 1 | 1])

    def get(self, i):
        i += self.N
        return self.X[i]

    def add(self, i, x):
        i += self.N
        self.X[i] += x
        while i > 1:
            i >>= 1
            self.X[i] = self.func(self.X[i << 1], self.X[i << 1 | 1])

    def update(self, i, x):
        i += self.N
        self.X[i] = x
        while i > 1:
            i >>= 1
            self.X[i] = self.func(self.X[i << 1], self.X[i << 1 | 1])

    def query(self, l, r):
        l += self.N
        r += self.N
        vl = self.e
        vr = self.e
        while l < r:
            if l & 1:
                vl = self.func(vl, self.X[l])
                l += 1
            if r & 1:
                r -= 1
                vr = self.func(self.X[r], vr)
            l >>= 1
            r >>= 1
        return self.func(vl, vr)


class DFSTree:
    # cf: https://codeforces.com/blog/entry/68138
    def __init__(self, N: int):
        self.N = N
        self.edges = []

        self.E = [[] for _ in range(self.N)]
        # span-edge and back-edge (directed)
        self.span_edge = [[] for _ in range(self.N)]
        self.back_edge = [[] for _ in range(self.N)]

        self.ord = [-1] * self.N
        self.low = [-1] * self.N
        self.par = [None] * self.N

        self.is_art = [False] * self.N
        self.is_bridge = []

        self.built = False

    def add_edge(self, u: int, v: int) -> None:
        """add edge"""
        eid = len(self.edges)
        self.edges.append((u, v))
        self.E[u].append((v, eid))
        self.E[v].append((u, eid))
        self.is_bridge.append(False)

    def build(self) -> None:
        """build dfs tree"""
        cnt = 0
        self.tecc_id = [-1] * self.N
        self.tvcc_id = [-1] * len(self.edges)
        self.bcc_num = 0
        for i in range(self.N):
            if ~self.ord[i]:
                continue
            stack = [(i, -1, -1)]
            estack = []
            while stack:
                v, p, p_eid = stack.pop()
                if v < 0:
                    v = ~v
                    for d, i in self.span_edge[v]:
                        if d == p:
                            continue
                        self.low[v] = min(self.low[v], self.low[d])
                    if ~p and self.ord[p] <= self.low[v]:
                        while True:
                            eid = estack.pop()
                            self.tvcc_id[eid] = self.bcc_num
                            if eid == p_eid:
                                break
                        self.bcc_num += 1
                else:
                    if ~self.ord[v]:
                        continue
                    self.ord[v] = cnt
                    self.low[v] = cnt
                    cnt += 1
                    # p -> v is span-edge.
                    if ~p:
                        self.par[v] = (p, p_eid)
                        self.span_edge[p].append((v, p_eid))
                        estack.append(p_eid)
                    stack.append((~v, p, p_eid))

                    for d, eid in self.E[v][::-1]:
                        if eid == p_eid:
                            continue
                        if ~self.ord[d]:
                            # v -> d is back-edge since d is already visited.
                            self.back_edge[v].append((d, eid))
                            self.low[v] = min(self.low[v], self.ord[d])
                            estack.append(eid)
                            continue
                        stack.append((d, v, eid))

        self._search_bridge()
        self._search_articulation_points()
        self.built = True

    def bridges(self) -> List[Tuple[int, int]]:
        """return list of edges that are bridges"""
        assert self.built
        return [e for i, e in enumerate(self.edges) if self.is_bridge[i]]

    def articulation_points(self) -> List[int]:
        """return list of vertices that are articulation points"""
        assert self.built
        return [i for i in range(self.N) if self.is_art[i]]

    def two_edge_connected_components(self) -> List[List[int]]:
        assert self.built
        cnt = 0
        for i in range(self.N):
            if ~self.tecc_id[i]:
                continue
            stack = [i]
            while stack:
                v = stack.pop()
                if ~self.tecc_id[v]:
                    continue
                self.tecc_id[v] = cnt
                for d, eid in self.E[v]:
                    if ~self.tecc_id[d] or self.is_bridge[eid]:
                        continue
                    stack.append(d)
            cnt += 1
        ret = [[] for _ in range(cnt)]
        for i, tid in enumerate(self.tecc_id):
            assert ~tid
            ret[tid].append(i)
        return ret

    def biconnected_components(self) -> List[List[int]]:
        assert self.built
        ret = [set() for _ in range(self.bcc_num)]
        for eid, tid in enumerate(self.tvcc_id):
            assert ~tid
            u, v = self.edges[eid]
            ret[tid].add(u)
            ret[tid].add(v)
        ret = [list(s) for s in ret]
        for i in range(self.N):
            if not self.E[i]:
                ret.append([i])
        return ret

    def _search_bridge(self) -> None:
        for u in range(self.N):
            for v, i in self.span_edge[u]:
                # (u, v) is bridge if vertex u has child v
                # that does not have lowlink to pass over its parent
                self.is_bridge[i] = self.ord[u] < self.low[v]

    def _search_articulation_points(self) -> None:
        for u in range(self.N):
            if self.par[u] is None:
                self.is_art[u] = len(self.span_edge[u]) >= 2
            else:
                for v, _ in self.span_edge[u]:
                    if self.ord[u] <= self.low[v]:
                        self.is_art[u] = True
                        break


class HLD:
    # reference: https://codeforces.com/blog/entry/53170
    def __init__(self, N, E, root: int = 0):
        self.N = N
        self.E = E
        self.root = root

        self.D = [0] * self.N
        self.par = [-1] * self.N
        self.sz = [0] * self.N
        self.top = [0] * self.N

        self.ord = [None] * self.N

        self._dfs_sz()
        self._dfs_hld()

    def path_query_range(self, u: int, v: int, is_edge_query: bool = False) -> List[Tuple[int, int]]:
        """return list of [l, r) ranges that cover u-v path"""
        ret = []
        while True:
            if self.ord[u] > self.ord[v]:
                u, v = v, u
            if self.top[u] == self.top[v]:
                ret.append((self.ord[u] + is_edge_query, self.ord[v] + 1))
                return ret
            ret.append((self.ord[self.top[v]], self.ord[v] + 1))
            v = self.par[self.top[v]]

    def subtree_query_range(self, v: int, is_edge_query: bool = False) -> Tuple[int, int]:
        """return [l, r) range that cover vertices of subtree v"""
        return (self.ord[v] + is_edge_query, self.ord[v] + self.sz[v])

    def get_index(self, v: int) -> int:
        return self.ord[v]

    def lca(self, u, v):
        while True:
            if self.ord[u] > self.ord[v]:
                u, v = v, u
            if self.top[u] == self.top[v]:
                return u
            v = self.par[self.top[v]]

    def _dfs_sz(self):
        stack = [(self.root, -1)]
        while stack:
            v, p = stack.pop()
            if v < 0:
                v = ~v
                self.sz[v] = 1
                for i, dst in enumerate(self.E[v]):
                    if dst == p:
                        continue
                    self.sz[v] += self.sz[dst]
                    # v -> E[v][0] will be heavy path
                    if self.sz[self.E[v][0]] < self.sz[dst]:
                        self.E[v][0], self.E[v][i] = self.E[v][i], self.E[v][0]
            else:
                if ~p:
                    self.D[v] = self.D[p] + 1
                    self.par[v] = p
                # avoid first element of E[v] is parent of vertex v if v has some children
                if len(self.E[v]) >= 2 and self.E[v][0] == p:
                    self.E[v][0], self.E[v][1] = self.E[v][1], self.E[v][0]
                stack.append((~v, p))
                for dst in self.E[v]:
                    if dst == p:
                        continue
                    stack.append((dst, v))

    def _dfs_hld(self):
        stack = [(self.root, -1)]
        cnt = 0
        while stack:
            v, p = stack.pop()
            self.ord[v] = cnt
            cnt += 1
            heavy_path_idx = len(self.E[v]) - 1
            for i, dst in enumerate(self.E[v][::-1]):
                if dst == p:
                    continue
                # top[dst] is top[v] if v -> dst is heavy path otherwise dst itself
                self.top[dst] = self.top[v] if i == heavy_path_idx else dst
                stack.append((dst, v))


N, M, Q = readints()
dfs_tree = DFSTree(N)
E = [[] for _ in range(N)]
for _ in range(M):
    a, b = readints()
    a -= 1
    b -= 1
    E[a].append(b)
    E[b].append(a)
    dfs_tree.add_edge(a, b)
dfs_tree.build()
T = dfs_tree.two_edge_connected_components()
NE = [[] for _ in range(len(T))]
for u in range(N):
    for v in E[u]:
        s, t = dfs_tree.tecc_id[u], dfs_tree.tecc_id[v]
        if s == t:
            continue
        NE[s].append(t)
        NE[t].append(s)
NE = [list(set(ne)) for ne in NE]

solver = HLD(len(T), NE)
st = SegTree(len(T), max, (0, -1))
st.build([(0, i) for i in range(len(T))])
H = [[] for _ in range(len(T))]
for _ in range(Q):
    t, *q = readints()
    if t == 1:
        u, w = q
        u -= 1
        idx = solver.get_index(dfs_tree.tecc_id[u])
        heapq.heappush(H[idx], -w)
        v = -H[idx][0]
        st.update(idx, (v, idx))
    else:
        u, v = q
        u -= 1
        v -= 1
        s, t = dfs_tree.tecc_id[u], dfs_tree.tecc_id[v]
        e = (0, len(T))
        ma = e
        for l, r in solver.path_query_range(s, t):
            v = st.query(l, r)
            ma = max(ma, v)
        if ma == e:
            print(-1)
            continue
        max_value, idx = ma
        heapq.heappop(H[idx])
        nxt = -H[idx][0] if H[idx] else 0
        st.update(idx, (nxt, idx))
        print(max_value)

yukicoder

結果

テストケース

ソースコード