AtCoder Beginner Contest 417 D,E,F,G問題メモ

AtCoder Beginner Contest 417

D - Takahashi's Expectation

• 高橋くんは、これから $N$ 個のプレゼントをもらいます。
• 高橋くんにはテンションという非負整数のパラメータがあり、テンションはプレゼントをもらうごとに変動します。
• それぞれのプレゼントは価値 $P$ 、テンション上げ度 $A$ 、テンション下げ度 $B$ という $3$ つのパラメータをもち、これらのパラメータによって高橋くんのテンションは次のように変動します。
  • もらったプレゼントの価値 $P$ がテンションの値以上であるとき、高橋くんはプレゼントに喜び、テンションが $A$ だけ増加する。
  • もらったプレゼントの価値 $P$ がテンションの値より小さいとき、高橋くんはプレゼントにがっかりし、テンションが $B$ だけ減少する。ただし、高橋くんのテンションの値が $B$ より小さかった場合、高橋くんのテンションは $0$ になる。

• $i$ 番目 $(1\le i\le N)$ に受け取るプレゼントの価値は $P _ i$ 、テンション上げ度は $A _ i$ 、テンション下げ度は $B _ i$ です。
• $Q$ 個の質問が与えられるので、その全てに答えてください。
• $i$ 番目 $(1\le i\le Q)$ の質問では、非負整数 $X _ i$ が与えられるので次の質問に答えてください。
  • 高橋くんのテンションがはじめ $X _ i$ だったときの、$N$ 個のプレゼントをすべて受け取ったあとの高橋くんのテンションを求めよ。

• $1\le N\le10000$
• $1\le P _ i\le500\ (1\le i\le N)$
• $1\le A _ i\le500\ (1\le i\le N)$
• $1\le B _ i\le500\ (1\le i\le N)$
• $1\le Q\le5\times10 ^ 5$
• $0\le X _ i\le10 ^ 9\ (1\le i\le Q)$
• 入力はすべて整数

D問題にしては難しめに感じた。実装もやや重め。

初期値 $X_i$ の上限はでかいものの、$P_i$ が小さいので、 初期値がでかい場合は開始後しばらくテンションが下がり続ける、という点に注目する。
（テンションMAXな高橋くんがプレゼントをもらう度に落ち込んでくのを想像すると少し面白い）

「下がり続けるとは限らなくなる」のは、テンションが $500$（$P_i$ の制約上限）以下になったタイミングである。
$B_i$ の累積和を計算しておけば、各 $X_i$ が「いつ $500$ 以下になるか（または最後までならないか）」は 二分探索等で計算できる。

最後までならない場合は $X_i - \sum_j B_j$ が答え。

$k$ 回目の後に $500$ 以下になる場合は、その際のテンションが $t = \displaystyle X_i - \sum_{j=1}^{k}B_j$ になるというところまではすぐに計算できる。「$k+1$ 個目以降のプレゼントにより、最終的に $t$ からどう変化するか」がわかればよい。

これは、遷移を後ろから考えることで求められる。

• $\mathrm{DP}[i,j]:=i$ 個目のプレゼントの後にテンションが $j$ だったときの、最終的なテンション

初期値は、$\mathrm{DP}[N,j]=j$ である。
DPを $i=N-1,N-2,...,0$ の順に求めていく。
$j$ の範囲は、$0～1000$ を求めておくとよい。（$500$ 以下になった後でも、一時的に $P_i+A_i \le 1000$ までは上昇しうるので）

$\mathrm{DP}[k,t]$ が、それぞれのクエリに対する答えとなる。

DPは、2次元で全部持つよりは、$i$ の次元は省略し、 $i=k$ まで進んだタイミングで各クエリの $\mathrm{DP}[t]$ を求めていくと、メモリの節約になる。

Python3

from bisect import bisect
from collections import defaultdict


def solve(n, presents, q, xxx):
    sagepoyo = [0]
    for i in range(n):
        sagepoyo.append(sagepoyo[-1] + presents[i][2])

    ans = [-1] * q
    t = defaultdict(list)
    for qi, x in enumerate(xxx):
        i = bisect(sagepoyo, x - 1000)
        if i == n + 1:
            ans[qi] = x - sagepoyo[n]
        else:
            t[i].append(qi)

    dp = list(range(1001))
    for i in range(n - 1, -1, -1):
        if i + 1 in t:
            for qi in t[i + 1]:
                x = xxx[qi]
                x -= sagepoyo[i + 1]
                ans[qi] = dp[x]

        p, a, b = presents[i]
        ndp = [0] * 1001
        for j in range(1001):
            if j <= p:
                ndp[j] = dp[j + a]
            else:
                ndp[j] = dp[max(0, j - b)]
        dp = ndp

    if 0 in t:
        for qi in t[0]:
            x = xxx[qi]
            ans[qi] = dp[x]

    return ans


n = int(input())
presents = [tuple(map(int, input().split())) for _ in range(n)]
q = int(input())
xxx = [int(input()) for _ in range(q)]
ans = solve(n, presents, q, xxx)
print('\n'.join(map(str, ans)))

E - A Path in A Dictionary

• $N$ 頂点 $M$ 辺の単純連結無向グラフ $G$ が与えられます。
  • $G$ の頂点は頂点 $1$, 頂点 $2$, $\ldots$, 頂点 $N$ と番号付けられており、$i$ $(1\leq i\leq M)$ 本目の辺は頂点 $U_i$ と頂点 $V_i$ を結んでいます。
• $G$ における頂点 $X$ から頂点 $Y$ への単純パスのうち辞書順最小のものを求めてください。

• なお、本問題の制約下で条件をみたすようなものが必ず存在することが証明できます。
• $T$ 個のテストケースが与えられるので、それぞれについて答えを求めてください。

• $1\leq T\leq 500$
• $2\leq N\leq 1000$
• $N-1\leq M\leq \min\left( \frac{N(N-1)}{2},5\times 10^4\right)$
• $1\leq X,Y \leq N$
• $X\neq Y$
• $1\leq U_i\lt V_i \leq N$
• $i\neq j$ ならば $(U_i,V_i)\neq (U_j,V_j)$
• 与えられるグラフは連結である。
• $1$ つの入力における $N$ の総和は $1000$ 以下である。
• $1$ つの入力における $M$ の総和は $5\times 10^4$ 以下である。
• 入力はすべて整数

「$X$ を出発し、番号の小さい頂点を優先的に訪れ、同じ頂点を2度通らないように $Y$ に辿り着けたら終了するDFS」をおこなう。
解法のシンプルさと思いつきやすさという点では、D問題より簡単かも。

ただし、注意しないと計算量が $O(N!)$ になりTLEする。その証明を含めると、若干難易度が上がる。

①---（②～999 の密なグラフ）
｜
(1000)

X=1  Y=1000
純粋なDFSだと、密なグラフを様々な探索順で試してしまい、TLE

TLEを防ぐには、一度調べた頂点からは探索を広げないなど、何らかの枝刈りが必要となるが、、、

DFSの過程で、以下のことが分かったとする。

• $P=(P_1,P_2,...,P_k)$ という暫定パスで $P_k$ まで辿り着いた。（$P_1=X$）
• この後、DFSを続けた結果、$P_k$ からは $Y$ まで辿り着けなかった。

この時、以下のことが言えると枝刈りできて嬉しい。

• 以降、他の暫定パスで $P_k$ まで辿り着いたとき、$P_k$ からの探索は省略できる。

しかし、$P_k$ までの暫定パスによって $P_k$ 以降に使える頂点は変化するので、 ある暫定パスでは辿り着けなくても、他の暫定パスでは $Y$ に辿り着ける場合もある。

 ,----------,
①----②----③   Pk=③, X=①, Y=④ としたとき、①→②→③ だと辿り着けないが、
      ｜         ①→③だと、以降②→④とすると辿り着ける。
      ④

じゃあそのパスが辞書順最小になる可能性もあるんだから、枝刈りできないじゃん！と思うが、実は以下のことは示せる。

• $P_k$ を使うパスは「辞書順最小になり得ない」

略証

X→P2→...→Pl→...→Pk      ←暫定パスP
            ↓
            ○→...→Y

X→P2→...→Pl→(...A...)→Pk→              NG（Yに辿り着けない）
            ↓
            .............→Pk→(...B...)→Y  OK

X→P2→...→Pl→...→W→...→Pk→                NG
            ↓
            ...............→Pk→...→W→...→Y  OK

X→P2→...→Pl→ V →W→...→Pk→                NG
            ↓
            ...............→Pk→...→W→ V →Y  OK

よって、各頂点は高々1回ずつDFSで探索すれば十分とわかり、$O(N+M)$ で求められる。

Python3

import sys

sys.setrecursionlimit(1010)

try:
    import pypyjit

    pypyjit.set_param(inlining=0)
except:
    pass


def solve(n, m, x, y, edges):
    links = [[] for _ in range(n)]
    for u, v in edges:
        u -= 1
        v -= 1
        links[u].append(v)
        links[v].append(u)
    for i in range(n):
        links[i].sort()

    x -= 1
    y -= 1
    path = [x]
    used = [False] * n
    used[x] = True

    def dfs(u):
        if u == y:
            return True

        for v in links[u]:
            if used[v]:
                continue
            path.append(v)
            used[v] = True
            if dfs(v):
                return True
            path.pop()

        return False

    dfs(x)
    return [p + 1 for p in path]


t = int(input())
for _ in range(t):
    n, m, x, y = map(int, input().split())
    edges = [tuple(map(int, input().split())) for _ in range(m)]
    ans = solve(n, m, x, y, edges)
    print(*ans)

F - Random Gathering

• $N$ 個の皿が、皿 $1,$ 皿 $2,\ldots,$ 皿 $N$ の順に左から並んでいます。
• はじめ、皿 $i\ (1\le i\le N)$ には $A _ i$ 個の石が入っています。
• この皿たちに対して $M$ 回の操作を行います。
• $i$ 回目 $(1\le i\le M)$ の操作では、$2$ つの整数 $L _ i,R _ i$ が与えられ、次の操作を順に行います。
  • 皿 $L _ i,$ 皿 $L _ i+1,\ldots,$ 皿 $R _ i$ の $R _ i-L _ i+1$ 個の皿に入っている石をすべて皿の上から取り除く。
  • $L _ i$ 以上 $R _ i$ 以下の整数を一様ランダムに $1$ つ取り、それを $x$ とする。
  • 取り除いた石をすべて皿 $x$ に乗せる。

• $i=1,2,\ldots,N$ について、$M$ 回の操作がすべて終了したときに皿 $i$ に置かれている石の個数の期待値を${}\bmod998244353$ で求めてください。

• $1\le N\le2\times10 ^ 5$
• $1\le M\le2\times10 ^ 5$
• $0\le A _ i\lt998244353\ (1\le i\le N)$
• $1\le L _ i\le R _ i\le N\ (1\le i\le M)$
• 入力はすべて整数

各操作は「区間 $[L_i,R_i]$ の全てを、区間の平均値で上書きする」に置き換えられ、 区間上書き更新、区間和取得の遅延セグメント木があれば実装できる。

期待値の線形性と遅延セグメント木を知っていれば、解法も思いつきやすく、素直な問題。

Python3

from atcoder.lazysegtree import LazySegTree


def solve(n, m, aaa, pairs):
    MOD = 998244353

    def op(a, b):
        # (区間長, 和)
        return (a[0] + b[0], (a[1] + b[1]) % MOD)

    def mapping(a, b):
        # 一律update値
        if a == -1:
            return b
        return (b[0], b[0] * a % MOD)

    def composition(a, b):
        if a == -1:
            return b
        return a

    init = [(1, a) for a in aaa]
    lst = LazySegTree(op, (0, 0), mapping, composition, -1, init)

    for l, r in pairs:
        l -= 1
        _, s = lst.prod(l, r)
        s *= pow(r - l, -1, MOD)
        s %= MOD
        lst.apply(l, r, s)

    ans = [lst.get(i)[1] for i in range(n)]
    return ans


n, m = map(int, input().split())
aaa = list(map(int, input().split()))
pairs = [tuple(map(int, input().split())) for _ in range(m)]
ans = solve(n, m, aaa, pairs)
print(*ans)

G - Binary Cat

• 文字列 $S_0,S_1$ を $S_0=$ 0, $S_1=$ 1 と定義します。
• クエリが $Q$ 個与えられるので、順に処理してください。
• $i$ 番目 $(1\leq i\leq Q)$ のクエリでは、$3$ つの整数の組 $(L_i,R_i,X_i)$ が与えられます。
  • $S_{L_i}$ と $S_{R_i}$ をこの順に連結した文字列を $S_{i+1}$ とする。
  • その後、$S_{i+1}$ の $X_i$ 文字目を求める。
  • なお、$X_i$ が $S_{i+1}$ の長さ以下であることは保証される。

• $1\leq Q\leq 5\times 10^5$
• $0\leq L_i,R_i\leq i$
• $1\leq X_i\leq 10^{18}$
• $X_i$ は $S_{i+1}$ の長さ以下である。
• 入力はすべて整数

二分木、というわけではないが、 上に向かって既存頂点への2本の辺（$L,R$ の区別つき）を持つ頂点がどんどん追加されてく感じ。

      ⓪    ①
  ↗↗ L↖↗R ↖↖
 ③      ②      ④
  R↖ ↗L L↖ ↗R ↑
    ⑤       ⑥   ｜
     L↖  R  ↑   /
        ⑦~~~｜~~~
         R↖ ｜L
            ⑧

辺で結ばれた頂点の関係を「親子関係」と称することにする。「子」から「親」に向かって2本の辺が出ている。

クエリ番号と頂点番号がずれるのは気持ち悪いので、以降、クエリ番号の方をずらして考える。つまり、

• 頂点 $i$ が表す文字列は $S_i$
• 頂点 $i$（$i \ge 2$）は、$L_i$ と $R_i$ を親に持つ
• クエリでは $2 \le i \le Q+1$ につき $S_i$ の $X_i$ 文字目を答える

$X_i$ は0-indexedとしておく。

$W_i$ を、各 $S_i$ の文字列の長さとする。
$W_0=W_1=1$ であり、$i \ge 2$ に対しては2つの親の合計 $W_i=W_{L_i}+W_{R_i}$ として求められる。

$S_i$ の $X_i$ 文字目が $L_i,R_i$ のどちらに含まれるかは、すぐ特定できる。

• $X_i \lt W_{L_i}$ の時、答えは $S_{L_i}$ の $X_i$ 文字目
• $X_i \ge W_{L_i}$ の時、答えは $S_{R_i}$ の $X_i-W_{L_i}$ 文字目

このルールに従って $X_i$ 文字目が含まれる方の親をどんどん遡っていき、 最終的に⓪と①のどちらに辿り着きますか？がクエリに対する答えとなる。

このままでは、以下の2点でTLE,MLEとなる。

• $W_i$ は、$i$ の増加とともに指数的に増えていきうるので、愚直に持つとオーバーフローする
• $i$ から⓪か①までたどり着くのに最悪 $O(i)$ かかるため、全体 $O(Q^2)$ となる

クエリで聞かれうる $X_i \le 10^{18}$ を超える部分の文字列については、捨ててもクエリに答える分には問題ない。
ある $i$ で $S_i$ の $10^{18}$ 文字目より後になった文字は、 以降の $S_i$ を用いた $S_j$ でも $10^{18}$ 文字目より前に来ることはなく、クエリで問われる対象となり得ない。

よって、$W_i=\min(10^{18}, W_{L_i}+W_{R_i})$ と考えてよい。

$S_i$ もこれに合わせ、先頭 $10^{18}$ 文字のみを表すとする。

この言い換えは、オーバーフロー対策だけでなく、後述のダブリングで遡る回数の見積もりについても関わってくる。

木を遡るのを効率化するというと、LCAなどで使われるダブリングがある。
各頂点から親を1回、2回、4回、8回、、、遡った頂点を前計算しておくことで、 任意の回数、親を遡ることを $\log{N}$ でおこなえる。

• ダブリングを理解したい #C++ - Qiita

しかし今回の場合、あくまで“木っぽいグラフ”なので、親は1つでなく2つずつある。
どっちの親を遡った結果を保持するべきか、定まらないように見えるが、、、

以下のようにすることで、ダブリングっぽいことがおこなえるようになる。

• 2つの親のうち、「$S_i$ に含まれる長さが長い方の親を辿る」のを「主要ルート」とする。
  • 同じ場合はどちらでもよいが、とりあえず $R$ 側とする
• 主要ルートを遡った先の頂点をダブリングで前計算する
• 頂点番号とともに「それが $S_{i}$ の何文字目から始まる部分文字列で使われるか」も記録する

主要ルートは、「$W_i$ が大きい方」ではなく、あくまで（$10^{18}$ で上限を切った）$S_i$ に含まれる長さが基準である点に注意。
「$W_{L_i}=10^{18}-5$」「$W_{R_i}=10^{18}$」なら、$S_{R_i}$ は5文字分しか $S_i$ に含まれていないので、$L_i$ が主要ルートとなる。

前計算する値を以下のように定義する。

• $P_{i,k}:=i$ から $2^k$ 回、主要ルートを遡った頂点番号
• $B_{i,k}:=S_{P_{i,k}}$ が、$S_i$ の何文字目から始まるか
  • （※ $i$ から $2^k$ 回遡った頂点が存在する $k$ に対してのみ定義）

これらは $i$ の小さい方、同じ $i$ では $k$ の小さい方から $O(Q \log{Q})$ で計算できる。

• $P_{i,0}=$ 2つの親のうち主要ルートとした方の頂点番号
• $B_{i,0}=$ 主要ルートが $L_i$ なら $0$、$R_i$ なら $W_{L_i}$
• $P_{i,k}=P_{P_{i,k-1}, k-1}$（$k \ge 1$）
• $B_{i,k}=B_{i,k-1}+B_{P_{i,k-1},k-1}$（$k \ge 1$）

以上で構築した $P_{i,k},B_{i,k}$ を使って実際のクエリで親を遡ることを考える。
以下の2操作を1セットとする。
「逸脱しない」というのは、遡った先の文字列中に $S_i$ の $X_i$ 文字目が存在することを意味する。

• 主要ルートだけを、逸脱しない範囲でダブリングで一気に遡る
• 主要ルートでない辺を1回遡る

この時、主要ルートでない辺を辿る（＝このセットを繰り返す）回数は、$\log{\max(X)} \simeq 60$ 回に抑えられる。

略証

①
      X            10^18
|-----+----Si----------|

XがSiの前半（X<=M/2）にありながら「主要ルートでない」場合、
主要ルートはRi, XはLi側に含まれることになる。

|--Li-+--|-------Ri------...|

W[Li] < M のため、以降は Wi < M の場合に帰着される。


②
                X    10^18
|----------Si---+------|

XがSiの後半（M/2<X）にありながら「主要ルートでない」場合、
主要ルートはLi, XはRi側に含まれることになる。

|------Li-----|-+---Ri----...|

この時、X←X-W[Li] となるため、新しい X は M/2 未満となる。
・W[Ri]=M の場合、以降は①に帰着される。
・W[Ri]<M の場合、以降は Wi < M の場合に帰着される。

主要ルートをダブリングで遡るパートでは、$2^k$ 回遡って逸脱しないかどうかを $B_{i,k} \le X_i \lt B_{i,k}+W_{P_{i,k}}$（…★）を満たすかどうかで判定する。

• $i$ に対し★を満たす最大の $k$ を二分探索で（または大きい方から順に試して）特定する
  • 見つかった $k$ を $k_j$ とする
• $i←P_{i,k_j}, X_i←X_i-B_{i,k_j}$ とする
• 新たな $i$ から再び★を満たす最大の $k$ を探す。この時、$k$ の探索範囲は $0～k_j-1$ とする。繰り返す。

2巡目以降、$2^{k_j}$ 以上遡るのが無駄なことは直前のループで分かっているので、 それを利用して探索範囲を減らすことで、このパートは $O(\log{i})$ で完了する。

1セットの計算量は $O(\log{Q})$、これを最大 $\log{\max(X)}$ 回繰り返すので、1クエリ $O(\log{Q}\log{\max(X)})$ となる。

全体 $O(Q \log{Q} \log{\max(X)})$ で解ける。

Python3

import sys

input = sys.stdin.buffer.readline


def solve(q, queries):
    def bisect(dbl, x):
        l = 0
        r = len(dbl)
        while l + 1 < r:
            mid = (l + r) // 2
            u, ofs = dbl[mid]
            if ofs <= x < ofs + widths[u]:
                l = mid
            else:
                r = mid
        return l

    lefts = [-1] * (q + 2)
    rights = [-1] * (q + 2)
    heavy = [-1] * (q + 2)
    widths = [1] * (q + 2)
    doubling = [[], []]
    ans = []
    INF = 1 << 60
    for i in range(q):
        l, r, x = queries[i]
        i += 2
        lefts[i] = l
        rights[i] = r
        widths[i] = min(INF, widths[l] + widths[r])
        if widths[l] >= widths[r]:
            heavy[i] = 0
            v = l
            offset = 0
        else:
            heavy[i] = 1
            v = r
            offset = widths[l]
        dbl = [(v, offset)]
        for skip in range(60):
            if len(doubling[v]) <= skip:
                break
            v, ofs = doubling[v][skip]
            offset += ofs
            if offset > INF:
                break
            dbl.append((v, offset))
        doubling.append(dbl)

        v = i
        x -= 1
        while v > 1:
            lv = lefts[v]
            if heavy[v] == 0 and widths[lv] <= x:
                x -= widths[lv]
                v = rights[v]
                continue
            if heavy[v] == 1 and widths[lv] > x:
                v = lv
                continue
            dbl = doubling[v]
            j = bisect(dbl, x)
            v, ofs = dbl[j]
            x -= ofs
        ans.append(v)

    return ans


q = int(input())
queries = [tuple(map(int, input().split())) for _ in range(q)]
ans = solve(q, queries)
print('\n'.join(map(str, ans)))