#833231 (PyPy3) No.235 めぐるはめぐる (5)

提出ソース
結果

問題	No.235 めぐるはめぐる (5)
ユーザー	terasa
提出日時	2023-01-14 16:05:56
言語	PyPy3 (7.3.15)
結果	AC
実行時間	4,558 ms / 10,000 ms
コード長	8,717 bytes
コンパイル時間	224 ms
コンパイル使用メモリ	82,452 KB
実行使用メモリ	159,756 KB
最終ジャッジ日時	2024-06-07 17:28:31
合計ジャッジ時間	16,998 ms
ジャッジサーバーID （参考情報）	judge4 / judge5
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	4,558 ms 156,304 KB
testcase_01	AC	2,659 ms 159,756 KB
testcase_02	AC	3,774 ms 156,600 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()))


S = TypeVar('S')
F = TypeVar('F')


class LazySegTree:
    # reference: https://github.com/shakayami/ACL-for-python
    # reference: https://maspypy.com/segment-tree-%E3%81%AE%E3%81%8A%E5%8B%89%E5%BC%B72
    # reference: https://betrue12.hateblo.jp/entry/2020/09/22/194541
    def __init__(self, N: int, op: Callable[[S, S], S], e: S,
                 mapping: Callable[[F, S], S],
                 composition: Callable[[F, F], F], id_: F):
        """ 遅延セグメント木

        Args:
            N (int): 配列の長さ
            op (Callable[[S, S], S]): 区間取得に用いる演算
            e (S): 全てのaに対して op(a, e) = a が成り立つ単位元
            mapping (Callable[[F, S], S]): dataに作用させる関数
            composition (Callable[[F, F], F]): lazyに作用させる関数 f(g(x))
            id_ (F): 全てのaに対して mapping(id_, a) = a が成り立つ恒等写像

        Note:
            任意の x, y ∈ S, f, g ∈ F に対して、
            - f(op(x, y)) = op(f(x), f(y))
            - f(g(x)) = (g ∘ f)(x)
            であることが必要

            例) RMQ and RAQ
            - min(x, y) + a = min(x + a, y + a)
            - ((x + b) + a) = x + (a + b)
        """
        self.N = N
        self.op = op
        self.e = e
        self.mapping = mapping
        self.composition = composition
        self.id = id_

        self.K = (self.N - 1).bit_length()
        self.size = 1 << (self.K)

        self.data = [e] * (self.size << 1)
        self.lazy = [id_] * (self.size)

    def build(self, A: List[S]) -> None:
        for i in range(self.N):
            self.data[self.size + i] = A[i]
        for i in range(self.size - 1, 0, -1):
            self.data[i] = self.op(self.data[i << 1], self.data[i << 1 | 1])

    def _eval_at(self, i: int, f: F) -> None:
        self.data[i] = self.mapping(f, self.data[i])
        if i < self.size:
            self.lazy[i] = self.composition(f, self.lazy[i])

    def _propagate_at(self, i: int) -> None:
        self._eval_at(i << 1, self.lazy[i])
        self._eval_at(i << 1 | 1, self.lazy[i])
        self.lazy[i] = self.id

    def _propagate_above(self, i: int) -> None:
        H = i.bit_length() - 1
        for h in range(H, 0, -1):
            self._propagate_at(i >> h)

    def _recalc_at(self, i: int) -> None:
        self.data[i] = self.op(self.data[i << 1], self.data[i << 1 | 1])

    def _recalc_above(self, i: int) -> None:
        while i > 1:
            i >>= 1
            self._recalc_at(i)

    def set(self, i: int, x: S) -> None:
        i += self.size
        self._propagate_above(i)
        self.data[i] = x
        self._recalc_above(i)

    def get(self, i) -> S:
        i += self.size
        self._propagate_above(i)
        return self.data[i]

    def prod(self, l: int, r: int) -> S:
        assert 0 <= l and l <= r and r <= self.N
        if l == r:
            return self.e
        l += self.size
        r += self.size
        self._propagate_above(l // (l & -l))
        self._propagate_above(r // (r & -r) - 1)
        vl = self.e
        vr = self.e
        while l < r:
            if l & 1:
                vl = self.op(vl, self.data[l])
                l += 1
            if r & 1:
                r -= 1
                vr = self.op(self.data[r], vr)
            l >>= 1
            r >>= 1
        return self.op(vl, vr)

    def all_prod(self) -> S:
        return self.data[1]

    def apply(self, l: int, r: int, f: F) -> None:
        assert 0 <= l and l <= r and r <= self.N
        if l == r:
            return
        l += self.size
        r += self.size
        l0 = l // (l & -l)
        r0 = r // (r & -r) - 1
        self._propagate_above(l0)
        self._propagate_above(r0)
        while l < r:
            if l & 1:
                self._eval_at(l, f)
                l += 1
            if r & 1:
                r -= 1
                self._eval_at(r, f)
            l >>= 1
            r >>= 1
        self._recalc_above(l0)
        self._recalc_above(r0)


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))


def op(a, b):
    a0, a1 = a >> 32, a & msk
    b0, b1 = b >> 32, b & msk
    x, y = (a0 + b0) % mod, (a1 + b1) % mod
    return x << 32 | y


def mapping(f, x):
    x0, x1 = x >> 32, x & msk
    x, y = (x0 + f * x1) % mod, x1
    return x << 32 | y


def composition(f, g):
    return (f + g) % mod


N = int(input())
A = readlist()
C = readlist()
E = [[] for _ in range(N)]
for _ in range(N - 1):
    a, b = readints()
    a -= 1
    b -= 1
    E[a].append(b)
    E[b].append(a)
mod = 10 ** 9 + 7
msk = (1 << 32) - 1

solver = HLD(N, E)
lst = LazySegTree(N, op, 0, mapping, composition, 0)
B = [None] * N
for i in range(N):
    B[solver.get_index(i)] = A[i] << 32 | C[i]
lst.build(B)

Q = int(input())
for _ in range(Q):
    t, *q = readints()
    if t == 0:
        x, y, z = q
        x -= 1
        y -= 1
        for l, r in solver.path_query_range(x, y):
            lst.apply(l, r, z)
    else:
        x, y = q
        x -= 1
        y -= 1
        ans = 0
        for l, r in solver.path_query_range(x, y):
            ans += lst.prod(l, r) >> 32
            ans %= mod
        print(ans)
yukicoder

結果

テストケース

ソースコード
