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

提出ソース

結果

問題	No.529 帰省ラッシュ
ユーザー	terasa
提出日時	2023-01-14 12:47:28
言語	PyPy3 (7.3.15)
結果	AC
実行時間	2,016 ms / 4,500 ms
コード長	11,186 bytes
コンパイル時間	2,339 ms
コンパイル使用メモリ	85,892 KB
実行使用メモリ	204,520 KB
最終ジャッジ日時	2023-08-26 19:08:55
合計ジャッジ時間	25,014 ms
ジャッジサーバーID （参考情報）	judge13 / judge14

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	149 ms 79,972 KB
testcase_01	AC	150 ms 80,200 KB
testcase_02	AC	152 ms 79,992 KB
testcase_03	AC	155 ms 80,036 KB
testcase_04	AC	305 ms 85,592 KB
testcase_05	AC	292 ms 85,388 KB
testcase_06	AC	307 ms 86,024 KB
testcase_07	AC	309 ms 85,880 KB
testcase_08	AC	1,746 ms 190,836 KB
testcase_09	AC	1,609 ms 184,884 KB
testcase_10	AC	1,699 ms 187,352 KB
testcase_11	AC	1,704 ms 189,384 KB
testcase_12	AC	1,280 ms 161,276 KB
testcase_13	AC	1,071 ms 204,520 KB
testcase_14	AC	1,174 ms 199,084 KB
testcase_15	AC	2,016 ms 195,692 KB
testcase_16	AC	2,005 ms 196,472 KB
testcase_17	AC	1,805 ms 203,832 KB
testcase_18	AC	1,622 ms 201,464 KB
testcase_19	AC	1,605 ms 195,736 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

結果

テストケース

ソースコード