AtCoder Beginner Contest 351 G問題メモ

AtCoder Beginner Contest 351

G - Hash on Tree

問題

頂点1を根とした $N$ 頂点の根付き木が与えられる
各頂点には $A_1,A_2,...,A_N$ の値が書き込まれている
頂点 $v$ に対し、 $f(v)$ を以下で定義する
- $v$ が葉の時、 $f(v)=A_v$
- それ以外の時、 $v$ の子を $c_1,c_2,...,c_k$ として、 $f(v)=f(c_1) \times f(c_2) \times ... \times f(c_k) + A_v$
$Q$ 回のクエリに答えよ
- $i$ 回目のクエリでは、頂点 $v_i$ に書き込まれた値 $A_{vi}$ を $x_i$ に更新する
- 各クエリにおいて、更新後の $f(1) \mod{998244353}$ を出力する
$2 \le N \le 2 \times 10^5$
$1 \le Q \le 2 \times 10^5$
$0 \le A_i,v_i \lt 998244353$

解法

1回のクエリの影響が大きくて、更新した頂点から根に至るまでの全ての祖先の $f(v)$ が変わる。
平衡二分木のように高さがそこまで大きくなければよいが、そうとも限らない。

更新方法をよく観察すると、もし「更新頂点からの根に至るまでのパス」がセグメント木に乗っていて、それ以外の頂点が更新されないと仮定するならば、効率的に更新できる。

      ❶         f(1) = f(2)*f(4) * f(3) + A1
    ／｜＼              ~~~~~~~~~
   ② ❸ ④      f(3) = f(7)      * f(6) + A3
  /   /\                ~~~~
 ⑤  ❻⑦        f(6) =     1     * f(10)+ A7
 /\  ｜
⑧⑨ ❿          f(10)=     0            + A10
    
                       └────┘       └─┘
                            Xi              Yi
                    更新されないと仮定

$X_i x + Y_i$ という一次関数の合成はモノイドの条件を満たすので、 $(X_i,Y_i)$ をセグメント木に載せることで、❶❸❻❿のいずれかの更新に関しては $O(\log{N})$ でおこなえる。

実際にはそれ以外の頂点に対する更新もあるわけで、それをどうするか？

HL分解

絶対に初心者でもわかるHL分解/Heavy-Light-Decomposition #AtCoder - Qiita

他の頂点も、本流（①③⑥⑩）から外れたところから、再帰的に同じようにパス化してしまえばよい。
子が複数ある頂点からは、最も頂点数の多い子を残すようにパスに分解していく。
すると、どの頂点からも、根に遡るまでに経由するパス（下記[②⑤⑧]や[①③⑥⑩]など）の個数は $O(\log{N})$ に抑えられることが保証される。

      ①
    ／｜＼
   ② ③ ④
  /   /\
 ⑤  ⑥⑦
 /\  ｜
⑧⑨ ⑩

[①③⑥⑩]
[②⑤⑧]
[④]
[⑦]
[⑨]

⑨を更新するとき、

[⑨]のセグメント木で、 $Y_9$ を更新（長さ1なのでほぼ無意味だが）
→先頭の $f(9)$ を取得
$f(9)$ をもって [②⑤⑧] のセグメント木の $X_5$ を更新
→先頭の $f(2)$ を取得
$f(2)$ をもって [①③⑥⑩] のセグメント木の $X_1$ を更新
→先頭の①の $f(1)$ を取得

とすると、更新が発生する箇所全てを辿ることができる。

辿るパス数が $O(\log{N})$ で、各パス（セグ木）の更新も $O(\log{N})$ なので、1クエリ $O(\log^2{N})$ で更新できることになる。

0への対応

今回の場合、問題となるのが、 $f(i) \equiv 0 \mod{998244353}$ になり得ることである。

$X_1=f(2) \times f(4) \times ... いっぱい$ であるときに、 $f(2)$ の更新を $X_1$ に反映しようとすると、

$X_1 / (更新前のf(2)) \times (更新後のf(2))$

とするのが手っ取り早いが、 $更新前のf(2)=0$ だとゼロ除算となり、それができない。

ゼロは分けて管理する。 $X_i$ の情報を $W_i,Z_i$ に分けて、

$W_i=$ かけあわせる $f(c)$ のうち、ゼロでないものだけの積
$Z_i=$ かけあわせる $f(c)$ のうち、ゼロであるものの個数

とすると、 $Z_i=0$ のとき $X_i=W_i$ 、 $Z_i \gt 0$ のとき $X_i=0$ で、更新も問題なくおこなえる。

よって、 $(W_i,Z_i,Y_i)$ の3つの値を、HL分解後のセグメント木に載せるとよい。

計算量は $O(N+Q\log^2{N})$ 。

（※分けて管理する代わりに、各頂点 $v$ に対して「 $v$ の子からなる、1点更新・区間積取得のセグメント木」を作ってもよい）

Python3

        
          
          
              
              import os
from typing import List
import sys
 
import numpy as np
 
 
def solve(n: int, vvv: np.ndarray, xxx: np.ndarray, data: np.ndarray,
          position: np.ndarray, roots: np.ndarray, parents: np.ndarray):
    MOD = 998244353
 
    # == SegmentTree for Numba ==
    # Numbaで使うにあたり、いくつかのアドホックな使用方法に注意する。
    # ・sgt_dataは2次元配列 (index, 1要素当たりのデータ数) で、Numbaで扱える型。型は全て揃える必要がある。
    # ・関数定義
    #     sgt_operation(i,j,k):  data[k] に data[i]*data[j] を上書き更新。i と k が同じindexでも正常に動くようにする。
    #     sgt_initialize(i): data[i] を単位元に戻す
    # ・基本的な使用法
    #   ・準備: sgt_n, sgt_features, sgt_operation, sgt_initialize を定義する
    #   ・初期化: sgt_data[sgt_offset:] に初期値を埋め、sgt_build() を呼ぶと、構築される。
    #   ・更新: sgt_data[sgt_offset + i] を更新し、sgt_upstream(i) を呼ぶ
    #   ・取得: sgt_range_query(l, r)
 
    sgt_n = n  # CUSTOMIZABLE セグメント木の N
    sgt_features = 3  # CUSTOMIZABLE セグメント木 1要素当たりのデータ数
    sgt_depth = 0
    _n = sgt_n
    while _n:
        _n >>= 1
        sgt_depth += 1
    sgt_offset = 1 << sgt_depth
    sgt_data = np.zeros(((sgt_offset << 1) + 2, sgt_features), np.int64)  # 末尾2つは range_query 時の一時変数用
 
    # CUSTOMIZABLE
    def sgt_operation(i, j, k):
        f2 = sgt_data[j, 2]
        if sgt_data[i, 2] > 0:
            f0 = sgt_data[i, 0]
            f1 = 0
        else:
            f0 = (sgt_data[j, 0] * sgt_data[i, 1] + sgt_data[i, 0]) % MOD
            if sgt_data[j, 2] > 0:
                f1 = 0
            else:
                f1 = sgt_data[i, 1] * sgt_data[j, 1] % MOD
        sgt_data[k, 0] = f0
        sgt_data[k, 1] = f1
        sgt_data[k, 2] = f2
 
    def sgt_initialize(i):
        sgt_data[i, 0] = 0
        sgt_data[i, 1] = 1
        sgt_data[i, 2] = 0
 
    # /CUSTOMIZABLE
 
    def sgt_build():
        """ 初期値 sgt_data[sgt_offset:] を埋めた後に呼び、sgt_data[:sgt_offset] を構築 """
        for i in range(sgt_offset - 1, 0, -1):
            lch = i << 1
            rch = lch | 1
            sgt_operation(lch, rch, i)
 
    def sgt_upstream(i):
        """ Aiを更新後、その情報を上へ伝播（更新は sgt_data を直接上書き。sgt_offset 分だけずらすのに注意） """
        i += sgt_offset
        while i > 1:
            i >>= 1
            lch = i << 1
            rch = lch | 1
            sgt_operation(lch, rch, i)
 
    def sgt_range_query(l, r):
        l += sgt_offset
        r += sgt_offset
 
        sml = -2
        smr = -1
        sgt_initialize(sml)
        sgt_initialize(smr)
 
        while l < r:
            if (l & 1) == 1:
                sgt_operation(sml, l, sml)
                l += 1
            if (r & 1) == 1:
                r -= 1
                sgt_operation(r, smr, smr)
            l >>= 1
            r >>= 1
        sgt_operation(sml, smr, sml)
        return sgt_data[sml]
 
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a > 0:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret
 
    sgt_data[sgt_offset:sgt_offset + n, :] = data
    for i in range(sgt_offset + n, sgt_offset + sgt_offset):
        sgt_initialize(i)
 
    sgt_build()
 
    q = len(vvv)
    result = np.zeros(q, np.int64)
    for qi in range(q):
        v = vvv[qi]
        x = xxx[qi]
        i = position[v]
        _, c1, c2 = sgt_data[sgt_offset + i]
        c0 = x
 
        while True:
            l = roots[i]
            e0, e1, e2 = sgt_range_query(l, sgt_offset)  # 現在のパス成分の根の値
            sgt_data[sgt_offset + i, 0] = c0
            sgt_data[sgt_offset + i, 1] = c1
            sgt_data[sgt_offset + i, 2] = c2
            sgt_upstream(i)
 
            if l == 0:
                break
 
            f0, f1, f2 = sgt_range_query(l, sgt_offset)  # 更新後のパス成分の根の値
            i = parents[l]
            c0, c1, c2 = sgt_data[sgt_offset + i]
 
            if e0 == 0:
                c2 -= 1
            else:
                c1 *= mod_pow(e0, MOD - 2, MOD)
                c1 %= MOD
            if f0 == 0:
                c2 += 1
            else:
                c1 *= f0
                c1 %= MOD
 
        result[qi] = sgt_data[1, 0]
 
    return result
 
 
SIGNATURE = '(i8,i8[:],i8[:],i8[:,:],i8[:],i8[:],i8[:])'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC
 
    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()
 
if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit
 
    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)
 
 
def heavy_light_decomposition(n: int, children: List[List[int]], root: int = 0):
    weights = [-1] * n
    q = [root]
    while q:
        u = q[-1]
        if weights[u] == -1:
            weights[u] = -2
            q.extend(children[u])
        else:
            q.pop()
            weights[u] = 1 + sum(weights[v] for v in children[u])
 
    q = [root]
    progress = [0] * n
    visited_order = []
    position = [0] * n
    current_subroot = root
    subroots = [0] * n
    while q:
        u = q[-1]
        if progress[u] == 0:
            children[u].sort(key=weights.__getitem__, reverse=True)
            position[u] = len(visited_order)
            visited_order.append(u)
            subroots[u] = current_subroot
        if progress[u] >= len(children[u]):
            q.pop()
            continue
        v = children[u][progress[u]]
        if progress[u] > 0:
            current_subroot = v
        progress[u] += 1
        q.append(v)
 
    return visited_order, position, subroots
 
 
input = sys.stdin.buffer.readline
 
n, q = map(int, input().split())
ppp = list(map(int, input().split()))
aaa = list(map(int, input().split()))
children = [[] for _ in range(n)]
parent = [-1] * n
for i, p in enumerate(ppp, start=1):
    p -= 1
    children[p].append(i)
    parent[i] = p
MOD = 998244353
 
hld_order, hld_position, hld_roots = heavy_light_decomposition(n, children, 0)
 
# Ai, children, parent の各情報を、hld_order の順に並べ直す
new_aaa = [0] * n
new_children = [[]] * n
new_parent = [0] * n
new_roots = [0] * n
for i in range(n):
    hpi = hld_position[i]
    new_aaa[hpi] = aaa[i]
    new_children[hpi] = [hld_position[c] for c in children[i]]
    new_parent[hpi] = hld_position[parent[i]]
    new_roots[hpi] = hld_position[hld_roots[i]]
aaa = new_aaa
children = new_children
parent = new_parent
hld_roots = new_roots
 
children_products = [1] * n  # 「0でない子」のみの f(c) を掛け合わせた値（Ai を足す前）、葉は 0
children_products2 = [1] * n  # HLDで繋がったままの子「以外」の「0でない子」のみの f(c) を掛け合わせた値（Ai を足す前）
# 葉は 0、子が1つのみの頂点は 1
zero_count = [0] * n  # 0である子の個数
zero_count2 = [0] * n  # HLDで繋がったままの子「以外」で 0 である子の個数
 
init = []
for i in range(n - 1, -1, -1):
    if len(children[i]) == 0:
        children_products[i] = 0
        children_products2[i] = 0
    else:
 
        first_child = children[i][0]
        for j in children[i]:
            zc = zero_count[j]
            if zc > 0:
                cp = aaa[j]
            else:
                cp = children_products[j] + aaa[j]
                cp %= MOD
            if cp == 0:
                zero_count[i] += 1
                if j != first_child:
                    zero_count2[i] += 1
            else:
                children_products[i] *= cp
                children_products[i] %= MOD
                if j != first_child:
                    children_products2[i] *= cp
                    children_products2[i] %= MOD
    a0 = aaa[i]
    a1 = children_products2[i]
    a2 = zero_count2[i]
    init.append((a0, a1, a2))
 
init.reverse()
 
vvv, xxx = zip(*(map(int, input().split()) for _ in range(q)))
vvv = np.array(vvv, np.int64) - 1
xxx = np.array(xxx, np.int64)
init = np.array(init, np.int64)
position = np.array(hld_position, np.int64)
roots = np.array(hld_roots, np.int64)
parents = np.array(parent, np.int64)
 
ans = solve(n, vvv, xxx, init, position, roots, parents)
 
print('\n'.join(map(str, ans)))

            

        
      

解法1.5

解法1で、セグメント木の作り方を敢えてアンバランスにすることにより、logを1つ落とすことができる、らしい。

HL decomposition+SegTreeのlogを1個消す - よすぽの日記

①--②--③--④--⑤--⑥--⑦--⑧    ←Heavy path
 |   |   |   |   |   |   |
 |   :   :   :   :   :   :        ←Light paths
 |`⑨-⑩-...
 |  `-⑪-...
 `-⑫-⑬-...
 
通常のセグメント木
  1  2  3  4  5  6  7  8
 |----------------------|
 |----------||----------|
 |----||----||----||----|
 |-||-||-||-||-||-||-||-|

アンバランス（例）
  1  2  3  4  5  6  7  8
 |----------------------|    ノードを左右に分割するとき、
 |----||----------------|    「ノード v から伸びる、Heavy path 以外の
 |-||-||-||-------------|      子の部分木サイズ」を Wv として
          |----||-------|    左右で Wv の合計がなるべく等しくなるように分割する
          |-||-||-||----|
                   |-||-|

ものすごくざっくりした理解で言うと、①～⑧がそれぞれの傍流のLight pathによって更新される際、「そこに至るまでに辿ってきた Light path の個数」や「辿ってきたセグメント木の深さ」（＝計算量）って、多く Light path の子孫を持ってる頂点ほど多くなりやすいはず。

そこで、「Light path を辿ってきた計算量」と「Heavy path 上のセグメント木の深さ」を均すことによって、最悪ケースが $O(\log{N})$ になる感じ。

ただ、配列で持てなくなるなど、軽い実装はしにくくなるので、 logまるまる1つ分の恩恵があるかというと、そうでもないらしい?（章冒頭の記事参照）。

解法2

公式Editorialでは、Static Top Tree というものを用いることで、 $O(N+Q\log{N})$ で解く方法が解説されている。

個々の頂点がバラバラに存在する状態から、元の木になるように統合していく操作は、5種類程度に分類できる。
元の木が平衡木や二分木でなくても、この「統合の過程」なら「なるべく平衡な二分木」で表すことができる。
（具体例は、公式の動画解説や「Static Top Treeの覚書」参照）

なるべく平衡な二分木の深さは $O(\log{N})$ （統合操作を表現するため元の $N$ よりノードが増えたり、完全には平衡にできないことで多少非効率な部分ができるにしても、4倍の定数倍がかかる程度）なので、クエリの更新は Static Top Tree を更新頂点から根まで辿る $O(\log{N})$ でおこなえる。