東京海上日動プログラミングコンテスト2024（AtCoder Beginner Contest 355）E,F,G問題メモ

東京海上日動プログラミングコンテスト2024（AtCoder Beginner Contest 355）

AIを使うにしても、一発かそこらでACコードを生成させるの、割とむずいと思うんだよな（使いこなせてない）。

E - Guess the Sum

• インタラクティブ問題
• ジャッジが隠し持っている長さ $2^N$ の数列 $A = (A_0, A_1, \dots, A_{2^N-1})$ の区間和 $\sum_{i=L}^{R}A_i$ を突き止めたい
• 質問は「セグメント木のノード的な区間」を指定して、その区間和を知ることができる
  • つまり、2つの整数 $i,j$ を決めて、$[j \times 2^i, (j+1) \times 2^i)$ で表せる区間
• 与えられた $N,L,R$ に対し「どのような $A$ であっても確実に答えを突き止められる回数の最小値」を $m$ とする
• $m$ 回以内の質問を行って $A_L+A_{L+1}+\dots+A_{R}$ を求めよ

• $1\leq N\leq 18$
• $0\leq L\leq R\leq 2^N-1$

回数の最小化が肝。

加算しか行えないなら「クエリ $[L,R]$ が与えられたときにセグメント木で辿るノード」を列挙すればよいが、 引き算を交えた方が回数が少なくなる場合もあるため、ややこしくなる。

公式Editorialにある解法は、グラフの問題にすること。直感的に気付きにくい。
だが、今回のセグメント木のような規則的な並びでなくても、区間和をいくつかの区間和の加減算で求めるのに汎用的に使える。

ある2つの区間和を加減算して、別の（連続）区間の区間和を求めるには、 2つの区間の「どちらかの端が揃っている」ことが必要となる。

|------|           |---------------|
     +                   -
       |-----|     |---|
     =                   =
|------------|         |-----------|

$[L,a),[a,b),[b,c),...,[y,z),[z,R)$ のように、次々と区間端が一致するような区間の集合があれば、 それぞれの区間和を足し引きして $[L,R)$ を求めることができる。

そして、この区間の数を最小化するというのは、$L$ から $R$ への最短経路問題に他ならない。

区間 $[a,b)$ という表記において、$a \gt b$ であっても許されるとし、
$a \lt b$ なら区間和 $[a,b)$ を答えに加算、$a \gt b$ なら $[b,a)$ を減算していけばよい。

最適性の証明

まず普通に「クエリ $[L,R]$ が与えられたときにセグメント木で辿るノード」を求めてから、スタックを利用して最適化する。

辿るノードは、左半分と右半分から、どちらも中央に向かって山を登っていくような配置になる。
対称的なので、左半分を考える。

|---------------------------------------------------------------|--...
|-------------------------------|---------------X---------------|--...
|---------------|---------------|---------------|---------------|--...
|-------|-------|-------|---X---|-------|-------|-------|-------|--...
|---|---|---|---|---|-X-|---|---|---|---|---|---|---|---|---|---|--...
|-|-|-|-|-|-|-|-|-|X|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|...

本来のセグメント木で辿るノード（X）は、中央に向かい ↗ の方向に並ぶ

左下にある方から、順にスタックに積んでいく。積む際は、そのノード区間を加算するのか減算するのかの情報を付与する。

■あるノードと前のノードの深さが、1だけ異なる場合

|---|-W-|    W を追加するとき、スタックの一番上が X で、
|-|X|---|    深さが1だけ異なる区間であった
   ↓
|---Y---|    更に1つ上のノード区間 Y から、
|---|---|    不要な Z を除くこととして回数は変わらない。
|Z|-|---|    Xをpopし、Z,Yの順にaddする。


■あるノードと前のノードの深さが、同じ場合
  ※本来、セグメント木で辿るノード（の左半分）の深さは全て異なるが、
    前述の処理または当処理の結果、同じになることがある。

|---Y---|---P---|    P を追加するとき、スタックの一番上が Y であり、
|---|---|---|---|    同じ深さだった
|Z|-|-|-|-|-|-|-|
       ↓
|-------Q-------|    Yをpopし、Qをaddする
|-------|-------|    回数は1回減る
|---|---|---|---|
|Z|-|-|-|-|-|-|-|


■その他の場合

そのままスタックにノードをaddする

これを繰り返していくと、左半分で辿るノードの個数が最小になる。

証明はできてない。
なんとなく、個数を減らすことのできる2つめの処理をするには、 スタックトップの階層を積極的に浅くして行く方がよいみたいなお気持ち。

右からも同じことを行い、最後に、 左と右のスタックトップを1つのノードにまとめられるならまとめたものが、質問すべきノード区間の集合となる。

Python3

def make_queries(l, r):
    i = 0
    l_ans = []
    r_ans = []
    while l < r:
        if r & 1:
            r -= 1
            r_ans.append((i, r))
        if l & 1:
            l_ans.append((i, l))
            l += 1
        l >>= 1
        r >>= 1
        i += 1
    return l_ans, r_ans


def optimize(queries):
    if len(queries) == 0:
        return queries
    queries.sort()
    result = [(-9, -9, -9)]
    for i, j in queries:
        if result[-1][1] == i:
            assert result[-1][2] == j ^ 1
            result.pop()
            result.append((1, i + 1, j >> 1))
        elif result[-1][1] + 1 == i:
            assert (result[-1][2] & 1) == (j & 1)
            _, pi, pj = result.pop()
            result.append((0, pi, pj ^ 1))
            result.append((1, i + 1, j >> 1))
        else:
            result.append((1, i, j))
    return result[1:]


n, l, r = map(int, input().split())
r += 1
left, right = make_queries(l, r)
left = optimize(left)
right = optimize(right)

if left and right and left[-1][1] == right[-1][1] and left[-1][2] ^ 1 == right[-1][2]:
    op1, li, lj = left.pop()
    op2, ri, rj = right.pop()
    assert op1 == op2 == 1
    left.append((1, li + 1, lj >> 1))

queries = left + right

ans = 0
for op, i, j in queries:
    print(f'? {i} {j}')
    res = int(input())
    if op == 1:
        ans += res
    else:
        ans -= res

ans %= 100
print(f'! {ans}')

再帰でセグメント木のノードを潜っていくのをイメージする。

現在のノード $v$ の表す区間範囲を $[a,b)$、$m=(a+b)/2$ とする。
$[a,b)$ が $[l,r)$ に完全に内包される場合はノード $v$ の採用が決定する。
共通部分を持たない場合は $v$ 以下の不採用が決定する。

それ以外でより深いノードに潜っていくとき、

        a             b
current |---======----|        ===: l,r の範囲
        a      m   m      b
next    |---===| + |==----|

通常のセグメント木では、このように分割して、2つの結果（使用するノード）を合成していけばよい。

一方、今回の場合は引き算を考慮して、$[a,b)$ を採用した上で、以下の結果を引くパターンも存在する。

        a      m   m      b
next2   |===---| + |--====|

2つを試し、より質問回数の少ない方を返していくと、実際に行うべき質問がわかる。

ただし実装を上手くしないと、反転させた結果をまた反転させて無限ループに陥ってしまったりする。

• 参考
  • 【競技プログラミング】ABC355【実況】 - YouTube

F - MST Query

• $N$ 頂点 $N-1$ 辺の重み付き無向グラフ $G$ が与えられ、はじめ、$G$ は木である
• $Q$ 個のクエリが与えられるので順に処理せよ
  • $G$ の頂点 $u_i,v_i$ の間に重み $w_i$ の辺を（永続的に）追加する
  • 毎クエリ後、$G$ の最小全域木に含まれる辺の重みの和を出力する

• $2\leq N\leq 2 \times 10^5$
• $1\leq Q\leq 2 \times 10^5$
• $1\leq a_i\lt b_i\leq N$
• $1\leq u_i\lt v_i\leq N$
• $1\leq c_i,w_i\leq 10$

制約がヒント。辺コストが10以下。

素直に考えると、既に構築された最小全域木があり、そこにクエリの辺 $(u,v,w)$ を追加する場合、 「（現在の最小全域木でⓤ-ⓥ間を結ぶパスの中でコスト最大の辺）$ \gt w$」なら辺をつなぎ替える、という操作が必要になる。

 ○--8--○--10--○         ○     ○--10--○
4｜     ｜6          →   4｜     ｜6
 ⓤ     ⓥ                 ⓤ--5--ⓥ
  ` - - '
    w=5 の辺を追加

最小全域木のコストを求めるアルゴリズム クラスカル法 を考える。
連結状態をUnion-Find等で管理し、 コストが小さい方から辺を見ていって（$(u,v,w)$ とする）、 その時点で頂点 $u$ と $v$ が未連結なら、その辺を採用してよい。全域木のコストに $w$ を加算する。

今回のように後からクエリで辺を追加する場合、 この「コストが小さい方から辺を見ていって」の順番を守ることができないのが問題となる。

だが、コストの種類数が少ないので、コスト毎に連結状態を管理できる。

• コスト1の辺だけで連結か？のUnion-Find $UF_1$
• コスト2以下の辺だけで連結か？のUnion-Find $UF_2$
• コスト3以下の辺だけで連結か？のUnion-Find $UF_3$
• ……

$(u,v,w)$ に対し、$UF_w$ において $u,v$ が既に連結なら、その辺は今の最小全域木のコストを改善しない。

非連結の場合、$w$ から1個ずつ増やして確認していって、$UF_x$ で初めて連結になった場合、 「現在の最小全域木でⓤ-ⓥ間を結ぶパスの中でコスト最大の辺」は、コスト $x$ だと特定できる。

$x-w$ だけ最小全域木のコストを改善できることになるので、差分更新していく。

計算量は、コスト上限を $C$ として、 $O(C(N+Q) \alpha(N))$。

「具体的にどの辺が入れ替わる」かは特定しなくても答えが求まるのは、面白い。

Python3

class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n

    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def find(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return False
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
        else:
            self.table[r1] = r2
            self.table[r2] += d1
        return True

    def get_size(self, x):
        return -self.table[self.root(x)]


n, q = map(int, input().split())
ufts = [UnionFind(n) for _ in range(11)]
tmp = 0
for _ in range(n - 1):
    a, b, c = map(int, input().split())
    a -= 1
    b -= 1
    for ui in range(c, 11):
        ufts[ui].unite(a, b)
    tmp += c

buf = []
for _ in range(q):
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1
    if ufts[w].find(u, v):
        buf.append(tmp)
        continue
    current_cost = 10
    for c in range(9, 0, -1):
        if not ufts[c].find(u, v):
            current_cost = c + 1
            break
    # print(f'{u=} {v=} {w=} {current_cost=}')
    tmp -= current_cost
    tmp += w
    buf.append(tmp)

    for c in range(w, current_cost):
        ufts[c].unite(u, v)

print('\n'.join(map(str, buf)))

公式Editorialでは、方針はほぼ同じものの、考え方がよりテクニカル。

コストそのものを1つ1つのUnion-Findに分散させ、 「1つのUnion-Findで連結成分が1つ減るたび、最小全域木のコストは1減る」と考えている。

少し実装が楽になる。

G - Baseball

• 長さ $N$ の数列 $P=(P_1,P_2,\dots,P_N)$ が与えられます。
• 高橋君と青木君が、数列 $P$ を使ってゲームを行います。
  • まず、高橋君が、 $1,2,\dots,N$ から $K$ 個の相異なる整数 $x_1,x_2,\dots,x_K$ を選びます。
  • 次に、青木君が、 $1,2,\dots,N$ から $1$ つの整数 $y$ を $P_y$ に比例する確率で選びます。
    • すなわち、整数 $y$ が選ばれる確率は $\dfrac{P_y}{\sum_{y'=1}^N P_{y'}}$
  • 青木君が $\displaystyle \min_{i=1,2,\dots,K} |x_i-y|$ のスコアを得ます。
• 高橋君は、青木君が得るスコアの期待値を最小化したいです。
• 高橋君が適切に $x_1,x_2,\dots,x_K$ を選んだときに、青木君が得るスコアの期待値の最小値を $\sum_{y'=1}^N P_{y'}$ 倍した値を求めてください。なお、出力すべき値は整数になることが証明できます。

• $1 \leq N \leq 5 \times 10^4$
• $1 \leq K \leq N$
• $0 \leq P_i \leq 10^5$
• $1 \leq \sum_{y'=1}^N P_{y'} \leq 10^5$

公式解説を読んでの自分用メモ。

Alien DP をする。その仮定で、DPの高速化にLARSCHというアルゴリズムも用いる。
いずれも、なんでそれで上手くいくのか正当性を証明するところまで理解するのは時間かけないと難しそう。。。

ま、使うだけなら理解してなくても使えるしいっか！！！¹⁾

• Alien DP
  • 解説 - AtCoder Beginner Contest 218 H - Red and Blue Lamps
• LARSCH
  • LARSCH Algorithm | Library
  • 簡易版 LARSCH Algorithm - noshi91のメモ

使うアルゴリズムも難しいが、初手の言い換えもなかなか気付きづらい。

具体的に、$P$ と、$X=(x_1,x_2,...,x_K)$ を仮定して考えてみる。$S=\sum_iP_i$ とする。

K=2
    i   1  2  3  4  5  6  7  8  9 10 11 12 13
    P   3  1  4  1  5  9  2  6  5  3  5  8  9
    X            x                    x

スコア  3  2  1  0  1  2  3  3  2  1  0  1  2
期待値  9  2  4  0  5 18  6 18 10  3  0  8 18   --計→  この場合の期待値 = 101
(のS倍)

$y$ が $x_i$ と $x_{i+1}$ の間の値に選ばれた時のスコアは、$x_i,x_{i+1}$ のみによって決まる。
スコアは $x_i$ と $x_{i+1}$ の位置を $0$ として、0,1,2,3,3,2,1,0 のように中央に向けて両端から1ずつ増加していく形となる。（$y < x_1$ と $x_K < y$ の範囲は少しだけ異なる）
期待値は、各スコアに $P_i$ を乗じたものとなる。

    i   ... 4  5  6  7  8  9 10 11 ...
    P   ... 1  5  9  2  6  5  3  5 ...
    X       x                    x
スコア      0  1  2  3  3  2  1  0
期待値      0  5 18  6 18 10  3  0

$x_i$ と $x_{i+1}$ の間の期待値の合計（上例での 5+18+6+18+10+3）は、 $x_i,x_{i+1}$ が与えられれば $O(1)$ で求めることができる。
「$P_i$ の累積和」「$i \times P_i$ の累積和」「$(N-i+1) \times P_i$ の累積和」を事前計算しておけばよい。

この期待値の合計を、グラフにおける頂点 $x_i→x_{i+1}$ への移動コスト $c(x_i,x_{i+1})$ と見なすと、 元の問題は、（始終点 $0,N+1$ を追加して）頂点 $0$ から $N+1$ まで、$K+1$ 回の移動で到達するときの最小コストを求める問題となる。

 i     0  1  2  3  4  5  6  7  8  9 10 11 12 13 14
 P        3  1  4  1  5  9  2  6  5  3  5  8  9  
期待値    9  2  4  0  5 18  6 18 10  3  0  8 18
       o---------->o------------------->o------->o
          c(0,4)         c(4,11)         c(11,14)
          9+2+4       5+18+6+18+10+3       8+18

これを素朴なDPで求めようとすると、「$DP[i,j]=$ いま頂点 $i$ にいて、$j$ 回移動しているときの最小コスト」みたいなのが必要になるが、状態 $O(NK)$、遷移 $O(N)$ で $O(N^2K)$ かかり、このままだとTLE。

ここまで言い換えれば、以下のサイトの解説がバチコリ当てはまる。

• Monge グラフ上の $d$-辺最短路長を計算するアルゴリズム | Kyopro Encyclopedia of Algorithms

Alien DP と呼ばれる手法を使って、上記のDPを $O(N^2 \log{\alpha})$ に高速化できる。$\alpha$ の内容は後述する。

• 使える前提条件
  • $d$ 回の移動での最小コストを $f(d)$ としたとき、$f$ は凸性を持っている

今回の場合、直感的に、いっぱい $x_i$ を指定できる（いっぱい移動する）方がスコア期待値も小さくなりそう。
よって(?)、下に凸なので使える。（正当な証明はコスト関数のMonge性を利用する。公式Editorial参照）

• まず1回の移動にペナルティ $\lambda$ を追加したものを、新しいコスト関数とする。$C(i,j)=c(i,j)+\lambda$
• そのコストの上で、「移動回数に制約を設けない最小コスト」を求める
  • $DP_{\lambda}[i]=$ ペナルティ $\lambda$ とした上で、$i$ に至るまでの最小コスト
• $F(\lambda) = DP_{\lambda}[N+1] - (K+1) \lambda$ とする
• すると、$f(K+1) = \max_{\lambda \in \mathbb{Z}} F(\lambda)$ となり、さらに $F$ は上に凸なので三分探索が使える
  • $\mathbb{Z}$ は整数全体の集合

他の方の解説などで図示されているが、$\lambda$ を固定した1回のDPは 「$y=f(x)$ のグラフと $y=-\lambda x+b$ の直線が接するときの切片 $b$」を求めていると解釈できる。
また $F(\lambda)$ は、「直線 $y=-\lambda x + b$ に、$x=K+1$ を代入したもの」、 つまり「接線の $x=K+1$ の時の $y$ 座標」ということができる。

下に凸である $y=f(x)$ の接線は、$y=f(x)$ より上に来ることはない。 接線上の点である $F(\lambda)$ も必ず $f(K+1)$ 以下になる。
唯一、ちょうど $\lambda$ が $x=K+1$ で接するような傾きの時に、$F(\lambda) = f(K+1)$ となる。
ここを探している。

計算量 $O(N^2 \log{\alpha})$ の $\alpha$ は、三分探索で $\lambda$ を調べる範囲の広さとなる。

結論から言うと、「$f$ における傾きの最小～最大」の範囲を調べればよい。
（$\lambda$ をペナルティとして扱い、符号が反転しているが、本質的には同じ）

前項のざっくり説明より、$\lambda$ の探索範囲は、$y=f(x)$ の $x=K+1$ における傾きを含んでいればよいことになる。
$f$ が凸なので、傾きは単調増加（上に凸なら単調減少）であり、 傾きの最小・最大は端の $f(1)$ や $f(N)$ 付近を考察すると求められる。

この問題では、$f(x)$ の最小は $K=N$ の時、$f(N+1)=0$。 最大は、$x_i$ を1個だけ指定した場合のコストを雑に見積もって、$f(2) \le NS$ あたり。

つまり $f$ は $NS$ から $0$ まで単調減少する。
「$f(2)=NS$ から、急激に $f(3)=0$ になる」みたいなケースが、あり得る最大の傾きとなる。 （実際にそうなることがあるかは分からんが、上限をざっくり見積もる上で）
$\lambda$ は $0～NS$ の範囲を探索すればよいとわかる。

Alien DPで $\lambda$ を固定しての1回毎のDPは、以下で表された。

• $DP_{\lambda}[i]=$ ペナルティ $\lambda$ とした上で、$i$ に至るまでの最小コスト

これは 状態 $O(N)$、遷移 $O(N)$ でまだ $O(N^2)$ かかってしまう。

これを $O(N)$（または簡易版で $O(N \log{N})$）に高速化するのが、LARSCH となる。

• 前提条件
  • コスト関数 $c(i,j)$ は、Monge性を持つ
  • DPの遷移は以下のように表現できる
    • $\displaystyle DP[i]=\min_{j=0,...,i-1}(DP[j]+c(j,i))$

Monge性を利用して、探索範囲を上手に絞っている、らしい（まだちゃんと中身を理解してない）。

Python3

from itertools import accumulate
from typing import Callable, List, TypeVar

T = TypeVar('T')


def simple_larsch(n: int, check: Callable[[int, int, List[T], List[int]], None],
                  identity: Callable[[], T] = lambda: 1 << 60, v0: T = 0):
    """
    https://noshi91.hatenablog.com/entry/2023/02/18/005856

    以下の制約で、NxN の下三角行列 A の、各行に対する最小値を求める。
      * A は Totally Monotone
      * i 行目の最小値が確定しないと i 列目にはアクセスできない（i=0は所与）
      * A は陽に持たなくてよく、後述のcheck関数で2要素比較できれば良い
        （したがって、check()次第で最小値以外にも適用可能。Aがその比較関数においてTMであれば。）

    典型的な用途:
        DP[i] = min(DP[j] + Cost(j,i)) （j=0..i-1）で更新されるDP[i]（i=0,...,N-1）を、
        Cost関数がMongeの場合に O(N logN) で求められる。

    check(i, k, mini, amin):
        A[i][k] < mini[i] なら、i列目の暫定最小値 mini[i] ← A[i][k], その列index amin[i] ← k に更新する関数
        A[i][k] は、DPの文脈では「DP[k] から i に遷移した場合に DP[i] 候補となる値」に相当。

    """

    # 初期化
    mini = [identity() for _ in range(n)]
    mini[0] = v0
    amin = [0] * n

    def solve(l, r):
        if r - l == 1:
            return
        m = (l + r) // 2
        for k in range(amin[l], amin[r] + 1):
            check(m, k, mini, amin)
        solve(l, m)
        for k in range(l + 1, m + 1):
            check(r, k, mini, amin)
        solve(m, r)

    solve(0, n - 1)

    return mini, amin


def get_check(n, acc_p, acc_ip, acc_pi, penalty):
    def check(i, k, mini, amin):
        val = mini[k] + penalty
        if k == 0:
            val += acc_pi[i - 1] - acc_p[i - 1] * (n - i + 1)
        elif i == n + 1:
            val += acc_ip[n] - acc_ip[k] - (acc_p[n] - acc_p[k]) * k
        else:
            m = (k + i) // 2
            val += acc_ip[m] - acc_ip[k] - (acc_p[m] - acc_p[k]) * k
            val += acc_pi[i - 1] - acc_pi[m] - (acc_p[i - 1] - acc_p[m]) * (n - i + 1)
        if val < mini[i]:
            mini[i] = val
            amin[i] = k

    return check


n, k = map(int, input().split())
ppp = list(map(int, input().split()))
s = sum(ppp)

acc_p = [0] + list(accumulate(ppp))
acc_ip = [0] + list(accumulate((i + 1) * p for i, p in enumerate(ppp)))
acc_pi = [0] + list(accumulate((n - i) * p for i, p in enumerate(ppp)))


l = -1
r = n * s
ans = 0
while l + 1 < r:
    m1 = (l + r) // 2
    m2 = m1 + 1
    mini1, _ = simple_larsch(n + 2, get_check(n, acc_p, acc_ip, acc_pi, m1))
    mini2, _ = simple_larsch(n + 2, get_check(n, acc_p, acc_ip, acc_pi, m2))
    res1 = mini1[n + 1] - m1 * (k + 1)
    res2 = mini2[n + 1] - m2 * (k + 1)
    if res1 < res2:
        l = m1
        ans = res2
    else:
        r = m1
        ans = res1

print(ans)