Ｓｋｙ株式会社プログラミングコンテスト2023（AtCoder Beginner Contest 329）E,G 問題メモ

Ｓｋｙ株式会社プログラミングコンテスト2023（AtCoder Beginner Contest 329）

E - Stamp

問題

長さ $N$ の英大文字からなる文字列 $S$ と、長さ $M$ の英大文字からなる文字列 $T$
長さ $N$ の文字列 $X=$ '###…#' に、以下の操作を好きなだけ繰り返して、 $S$ にすることができるか判定せよ
- $X$ の連続する長さ $M$ の好きな位置を、 $T$ で置き換える
$1 \le N \le 2 \times 10^5$
$1 \le M \le 5$

解法

$M$ が小さいので、「 $S$ の $i$ 文字目と合わせることが可能な $T$ のindexの集合」を持って、 $S$ の左から状態を管理する。

   12345
T: ABABC

S: A B A B A B C C A B C
   1 2 3 4
       1 2 3 4 5
           1 2
                 5
                   1 2
                   3 4 5

$S$ の $i$ 文字目と合わせられる $T$ のindexの集合を $D_i$ とする。（たとえば $D_3=\{1,3\}, D_4=\{2,4\}$ など）

$D_{i-1}$ がわかっているとき、 $D_i$ は、以下を合わせたものとなる。

1つ前から継続する
- $D_{i-1}$ に含まれる要素（ $k$ とする）について、 $S_i = T_{k+1}$ なら、 $k+1$
$i-1$ より左を作った後、 $i$ にスタンプの左端を合わせて押す
- $S_i=T_1$ なら $1$
$i$ より右を作った後、 $i-1$ にスタンプの右端を合わせて押す
- $M$ が $D_{i-1}$ に含まれる場合のみ、 $S_i = T_j$ であるような任意の $j$

で、途中で $D_i$ がなくなったらアウト。

なくならずに最後までいって、 $M$ が $D_N$ に含まれていれば成功となる。

Python3

G - Delivery on Tree

問題

$N$ 頂点の根付き二分木があり、根である頂点1に空のカゴがある
$M$ 個のボールがある。それぞれ頂点 $S_i$ に現在位置し、 $T_i$ に届けたい
カゴにボールを積み下ろしすることで、全てのボールを届けたい頂点に持っていきたい
各頂点では以下の3つができる
- カゴを隣接する頂点に移動させる
- $S_i$ でボール $i$ をカゴに積む
- $T_i$ でボール $i$ をカゴから下ろす
  - （つまり、 $S_i$ でも $T_i$ でもない頂点でボール $i$ を積み下ろしはできない）
カゴに一度に積めるボールは $K$ 個まで
各辺を1往復ずつする軌跡（オイラーツアー）のうち、全てのボールを届けたい頂点に持っていくことが可能な軌跡の個数を $\mod{998244353}$ で求めよ
$2 \le N \le 10^4$
$1 \le M \le 2 \times 10^5$
$1 \le K \le 10^3$
各頂点、子は最大2個まで

解法

比較的素直な解法だが、素直にやれば解けるという確信に至るまでの丁寧な考察と、実装力が問われる。

可能不可能の判定

まず、可能不可能を判定する。

LCA（最小共通祖先）を挟んで両側に $(S_i,T_i)$ があったら、 $S_i$ を先に訪れる必要がある。

  Ｃ     こんな(S,T)があったら、Cからは左を先に訪れる必要がある
  /\
 Ｓ○
   ｜
   Ｔ

これがボール間で矛盾すると答えは“0”。

矛盾せず、子が2個ある各頂点につき「左/右を先に訪れなければならない」「どちらでもよい」が決まったとする。

この「どちらでもよい」でどうするかが唯一の自由度となるが、その頂点を訪れたときに既にカゴに積まれていたボールの個数によっては、どちらか/両方が不可能な場合もある。

個数ごとに場合分けが必要そう。（制約も $NK=10^7$ といい感じだし）

各頂点ですべき操作は決まる

ある頂点でカゴにボールを積み下ろしする最適な順番は、無駄なボールを積まないことを意識するとある程度決まる。
仮に、先に訪れる子を左の子とする。

①カゴが親から $v$ に着く
②親方面に $S_i$ があって、 $T_i=v$ のボールがあったら（積んできてるはずなので）下ろす
③左の部分木に $T_i$ があって、 $S_i=v$ のボールを積む
④左の子に行く
⑤…（左の子探索中）…
⑥左の子から帰ってくる
⑦左の部分木に $S_i$ があって、 $T_i=v$ のボールを下ろす
⑧右の部分木に $T_i$ があって、 $S_i=v$ のボールを積む
⑨右の子に行く
⑩…（右の子探索中）…
⑪右の子から帰ってくる
⑫右の部分木に $S_i$ があって、 $T_i=v$ のボールを下ろす
⑬親方面に $T_i$ があって、 $S_i=v$ のボールを積む
⑭親方面に帰る

この操作と異なる手順でボールを積み下ろししても、悪くなることはあってもよくなることはない。
（※悪くなる＝ $v$ に来た時点でのカゴのボールの個数が同じ状況で、上記の操作では $K$ 個を超えないが、異なる操作では超えてしまう）
（※よくなる＝ $v$ に来た時点でのカゴのボールの個数が同じ状況で、上記の操作では $K$ 個を超えてしまうが、異なる操作では超えない）

なので、親・左・右別に、 $S_i=v,T_i=v$ となる個数、計6通りを数えておくと、操作を追えそう。

子が1個なら①～⑦⑬⑭、子がない葉頂点なら①②⑬⑭だけでよいが、いずれにしろ最適な操作は決まる。

木DP

$DP[i,j]=$ 頂点 $i$ に、カゴの中のボール $j$ 個で親から訪れたときに、取れるパターン数

$i,j$ に対し、①～⑭の操作をたどることで遷移を考える。
先ほどの操作の番号を、そのまま積み下ろしするボールの個数として扱うと、

$DP[v,k]=DP[l,k-②+③] \times DP[r,k-②+③+⑤-⑦+⑧]$ … (A)

となる。ただし、 $l,r$ は $v$ の左右の子、⑤は「左に行って帰ってきた際に変動するボールの個数（左の部分木内で積まれた数-下ろされた数）」とする。
⑤や⑩は、 $S_i$ に+1、 $T_i$ に-1をして、葉から根に向かって累積和を取っていくことで各頂点計算できる。

各操作の途中で $0$ 個を下回ったり、 $K$ 個を超えるような $k$ は無効なので、その場合は $DP[v,k]=0$ となる。
②,③,⑤,⑦,⑧,…,⑬ でカゴのボールの個数が変化するごとに逐一個数チェックし、途中一度でも超えるとアウト。

上記の遷移(A)は、子が1個の場合は $DP[l,k-②+③]$ のみとなる。
また、どちらを先に訪れても良い頂点なら、右頂点を先に訪れた場合も同様に計算して足し合わせる。

このように木DPをおこない、最終的に $DP[1,0]$ が答えとなる。 $O(NK)$

Python3

        
          
          
              
              class LcaDoubling:
    """
    links[v] = { u1, u2, ... }  (u:隣接頂点, 親は含まない)
    というグラフ情報から、ダブリングによるLCAを構築。
    任意の2頂点のLCAを取得できるようにする
    """
 
    def __init__(self, n, links, root=0):
        self.depths = [-1] * n
        prev_ancestors = self._init_dfs(n, links, root)
        self.ancestors = [prev_ancestors]
        max_depth = max(self.depths)
        d = 1
        while d < max_depth:
            next_ancestors = [prev_ancestors[p] for p in prev_ancestors]
            self.ancestors.append(next_ancestors)
            d <<= 1
            prev_ancestors = next_ancestors
 
    def _init_dfs(self, n, links, root):
        q = [(root, -1, 0)]
        direct_ancestors = [-1] * (n + 1)  # 頂点数より1個長くし、存在しないことを-1で表す。末尾(-1)要素は常に-1
        while q:
            v, p, dep = q.pop()
            direct_ancestors[v] = p
            self.depths[v] = dep
            q.extend((u, v, dep + 1) for u in links[v])
        return direct_ancestors
 
    def get_lca(self, u, v):
        du, dv = self.depths[u], self.depths[v]
        if du > dv:
            u, v = v, u
            du, dv = dv, du
        tu = u
        tv = self.upstream(v, dv - du)
        if u == tv:
            return u
        for k in range(du.bit_length() - 1, -1, -1):
            mu = self.ancestors[k][tu]
            mv = self.ancestors[k][tv]
            if mu != mv:
                tu = mu
                tv = mv
        lca = self.ancestors[0][tu]
        assert lca == self.ancestors[0][tv]
        return lca
 
    def upstream(self, v, k):
        i = 0
        while k:
            if k & 1:
                v = self.ancestors[i][v]
            k >>= 1
            i += 1
        return v
 
 
def solve(n, m, k, ppp):
    parents = [-1]
    children = [[] for _ in range(n)]
    for i, p in enumerate(ppp, start=1):
        p -= 1
        parents.append(p)
        children[p].append(i)
 
    lca = LcaDoubling(n, links=children)
    order_limit = [-1] * n  # 左右の子、どちらを先に訪れなければならないか 0:左 1:右 -1:制約無し
    p_load = [0] * n  # ここで積んで親方面へ持ってくボールの数
    p_drop = [0] * n  # 親方面から持ってきてここで降ろすボールの数
    l_load = [0] * n  # 左（1つめ）の子へ持ってく
    l_drop = [0] * n
    r_load = [0] * n
    r_drop = [0] * n
    diff = [0] * n  # 親に繋がる辺を入って出るときに差し引き増えてるボールの数
 
    for _ in range(m):
        s, t = map(int, input().split())
        s -= 1
        t -= 1
        c = lca.get_lca(s, t)
 
        diff[s] += 1
        diff[t] -= 1
 
        if s == c:
            if len(children[s]) == 1:
                l_load[s] += 1
            else:
                # t が c=s のどっちの子か特定
                v = children[c][0]
                d = lca.get_lca(t, v)
                if d == v:
                    l_load[s] += 1
                else:
                    r_load[s] += 1
            p_drop[t] += 1
        elif t == c:
            if len(children[t]) == 1:
                l_drop[t] += 1
            else:
                # s が c=t のどっちの子か特定
                v = children[c][0]
                d = lca.get_lca(s, v)
                if d == v:
                    l_drop[t] += 1
                else:
                    r_drop[t] += 1
            p_load[s] += 1
        else:
            p_load[s] += 1
            p_drop[t] += 1
            # c は子を2個持つこと確定。s,t がそれぞれ c のどっちの子か特定
            v = children[c][0]
            d = lca.get_lca(s, v)
            if d == v:
                ol = 0
            else:
                ol = 1
            if order_limit[c] == -1:
                order_limit[c] = ol
            elif order_limit[c] != ol:
                return 0
 
    MOD = 998244353
    q = [0]
    status = [0] * n
    dp = [[0] * (k + 1) for _ in range(n)]
    while q:
        v = q[-1]
        if status[v] == 0:
            status[v] = 1
            q.extend(children[v])
            continue
 
        q.pop()
        dpv = dp[v]
        if len(children[v]) == 0:
            pl = p_load[v]
            pd = p_drop[v]
            x_max = min(k, k + pd - pl)
            for x in range(pd, x_max + 1):
                dpv[x] = 1
        elif len(children[v]) == 1:
            c = children[v][0]
            d = diff[c]
            dpc = dp[c]
            pl = p_load[v]
            pd = p_drop[v]
            ll = l_load[v]
            ld = l_drop[v]
            x_min = max(0, pd, pd - ll - d, pd - ll - d + ld)
            x_max = min(k, k + pd - ll, k + pd - ll - d, k + pd - ll - d + ld - pl)
            for x in range(x_min, x_max + 1):
                dpv[x] = dpc[x - pd + ll]
        else:
            c0, c1 = children[v]
            pl = p_load[v]
            pd = p_drop[v]
            ll = l_load[v]
            ld = l_drop[v]
            rl = r_load[v]
            rd = r_drop[v]
            if order_limit[v] == 0:
                two_children_update(dpv, dp[c0], dp[c1], diff[c0], diff[c1], pl, pd, ll, ld, rl, rd, MOD)
            elif order_limit[v] == 1:
                two_children_update(dpv, dp[c1], dp[c0], diff[c1], diff[c0], pl, pd, rl, rd, ll, ld, MOD)
            else:
                two_children_update(dpv, dp[c0], dp[c1], diff[c0], diff[c1], pl, pd, ll, ld, rl, rd, MOD)
                two_children_update(dpv, dp[c1], dp[c0], diff[c1], diff[c0], pl, pd, rl, rd, ll, ld, MOD)
 
        if v > 0:
            p = parents[v]
            diff[p] += diff[v]
 
    return dp[0][0]
 
 
def two_children_update(dpv, dp0, dp1, d0, d1, pl, pd, ll, ld, rl, rd, MOD):
    x_min = max(0, pd, pd - ll - d0, pd - ll - d0 + ld, pd - ll - d0 + ld - rl - d1,
                pd - ll - d0 + ld - rl - d1 + rd)
    x_max = min(k, k + pd - ll, k + pd - ll - d0, k + pd - ll - d0 + ld - rl,
                k + pd - ll - d0 + ld - rl - d1, k + pd - ll - d0 + ld - rl - d1 + rd - pl)
    for x in range(x_min, x_max + 1):
        dpv[x] += dp0[x - pd + ll] * dp1[x - pd + ll + d0 - ld + rl]
        dpv[x] %= MOD
 
 
n, m, k = map(int, input().split())
ppp = list(map(int, input().split()))
 
ans = solve(n, m, k, ppp)
print(ans)

            

        
      