トヨタ自動車プログラミングコンテスト2024#7（AtCoder Beginner Contest 362）F,G問題メモ

トヨタ自動車プログラミングコンテスト2024#7（AtCoder Beginner Contest 362）

F - Perfect Matching on a Tree

• $N$ 頂点の木 $T$ が与えられます。$T$ の頂点には $1$ から $N$ の番号がついており、 $i\,(1\leq i \leq N-1)$ 番目の辺は頂点 $u_i$ と頂点 $v_i$ を双方向に結んでいます。
• $T$ を用いて、$N$ 頂点の完全グラフ $G$ を次のように定めます。
  • $G$ の頂点 $x$ と頂点 $y$ の間の辺の重み $w(x,y)$ を、$T$ における頂点 $x$ と頂点 $y$ の間の最短距離とする
• $G$ の最大重み最大マッチングを一つ求めてください。
  • すなわち、$\lfloor N/2 \rfloor$ 個の頂点のペアの集合 $M=\{(x_1,y_1),(x_2,y_2),\dots,(x_{\lfloor N/2 \rfloor},y_{\lfloor N/2 \rfloor})\}$ であって、各頂点 $1,2,\dots, N$ が $M$ に現れる回数がたかだか $1$ 回であるようなもののうち、 $\displaystyle \sum_{i=1}^{\lfloor N/2 \rfloor} w(x_i,y_i)$ が最大であるものを一つ求めてください。

• $2 \leq N \leq 2 \times 10^5$
• $1 \leq u_i \lt v_i \leq N$
• 入力されるグラフは木である

木の直径など、遠い頂点対から貪欲にマッチングさせていけば何となくよさそう？
その場合、最終的に木の真ん中らへんの頂点が残る。 言い換えると、各ペアを結んだパスは、必ず木の真ん中らへんを通るものとなる。

木の真ん中というと、木の重心が連想される。

• ツリーの重心分解 (木の重心分解) の図解 #競技プログラミング - Qiita

重心は、その頂点を根とすると全ての部分木のサイズが $N/2$ 以下になる頂点のことなので、 重心を必ず通るような頂点対をマッチングさせていけば、必ず $N/2$ ペアまで構築できる。
（$N$ が偶数で重心が1点の時は、残った1頂点と重心をペアとする）

この発想が正しいか確認する。

木の重心を通らないペア $P$ が答えの1つにあるとする。
この時必ず、木の重心で分解したときに、両頂点とも $P$ と異なる部分木に含まれるようなペアが存在する。
（存在しないと数が合わなくなるので）

◎--o,         ,-o    ○: 木の重心
     o-o-○-o-●-o    ◎: 木の重心を通らないペア
◎--o'    `-● `-o    ●: ◎と異なる部分木にある他のペア

なら、◎-● を結ぶ2組のペアにつなぎ替えて損しない。
よって、最適解の1つには、全てのペアが重心を経由しているものが存在する。

木の重心を求めるのは既によく知られた方法があり、適当な頂点から部分木のサイズを求める木DPをおこなえば分かる。 （上述のけんちょん氏の記事など参照）

重心からの各部分木に含まれる頂点が分かっているとき、これらをどうやってペアにするのがよいか？
適当にやると、頂点数が多い部分木が1つだけ残ってしまい、失敗しかねない。

全て $N/2$ 以下であることが分かっているので、 部分木の頂点リストをflattenし、$i$ と $i+N/2$ 番目をそれぞれペアにしていけば、楽に実装できる。

部分木1に含まれる頂点番号 [ 3, 7, 5, 8, ... ]
部分木2に含まれる頂点番号 [ 9, 3, 6, ... ]
部分木3に含まれる頂点番号 [ 2, 0, 1, ... ]
  ↓flatten
[ 3, 7, 5, 8, ..., 9, 3, 6, ..., 2, 0, 1, ..., 重心 ]
  ^                      ^
     ^                      ^        N/2 ずつ離れた頂点をペアにしていく
        ^                      ^

この発想が出ず、少し手間のかかる実装をしてしまった。
まぁ、基本方針が合ってたのでよし。

Python3

from collections import defaultdict
from heapq import heapify, heappop, heappush
from typing import Iterable, List, Set, Optional, Dict, Tuple


def get_centroid(n: int, links: List[Iterable[int]], s: int, removed: Optional[Set[int]] = None) \
        -> Tuple[List[int], Dict[int, int]]:
    """
    木の重心を見つける。
    重心は1つまたは2つあるため、リストで返す。また、s を根とした際の部分木のサイズも返す。

    繰り返し重心分解する場合にも適用可能。
    removedには既に重心として確定された頂点集合を与え、そこにある頂点は探索されないようになる。
    n には木全体ではなく、探索範囲の頂点数を与える。（1回前に探索した部分木のサイズから計算できる）

    :param n: 木の（探索範囲の）頂点数
    :param links: 隣接リスト
    :param s: 探索開始頂点
    :param removed: 探索しない頂点集合
    :return:
    """
    if removed is None:
        removed = set()

    parents = {s: -1}
    subtree_sizes = {s: n}
    status = defaultdict(int)
    q = [s]
    limit = n // 2
    centroid = []
    while q:
        u = q[-1]
        if status[u] == 0:
            status[u] = 1
            for v in links[u]:
                if parents[u] == v:
                    continue
                if v in removed:
                    continue
                parents[v] = u
                q.append(v)
        else:
            status[u] = 2
            q.pop()
            my = 1
            is_centroid = True
            for v in links[u]:
                if status[v] != 2:
                    continue
                if v in removed:
                    continue
                scv = subtree_sizes[v]
                my += scv
                if scv > limit:
                    is_centroid = False
            subtree_sizes[u] = my
            if n - my > limit:
                is_centroid = False
            if is_centroid:
                centroid.append(u)
    return centroid, subtree_sizes


n = int(input())
links = [set() for _ in range(n)]
for _ in range(n - 1):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    links[u].add(v)
    links[v].add(u)

c, ss = get_centroid(n, links, 0)
if len(c) == 2:
    _, vvv1 = get_centroid(n // 2, links, c[0], {c[1]})
    _, vvv2 = get_centroid(n // 2, links, c[1], {c[0]})
    assert len(vvv1) == len(vvv2)
    vvv1 = list(vvv1.keys())
    vvv2 = list(vvv2.keys())
    for v1, v2 in zip(vvv1, vvv2):
        print(f'{v1 + 1} {v2 + 1}')
else:
    c = c[0]
    vvv = []
    hq = []
    for u in links[c]:
        if ss[c] > ss[u]:
            _, vvv1 = get_centroid(ss[u], links, u, {c})
        else:
            _, vvv1 = get_centroid(n - ss[c], links, u, {c})
        vvv.append(set(vvv1.keys()))
        hq.append((-len(vvv1), len(vvv) - 1))
    heapify(hq)
    while len(hq) > 1:
        l1, vi1 = heappop(hq)
        l2, vi2 = heappop(hq)
        v1 = vvv[vi1].pop()
        v2 = vvv[vi2].pop()
        print(f'{v1 + 1} {v2 + 1}')
        if l1 < -1:
            heappush(hq, (l1 + 1, vi1))
        if l2 < -1:
            heappush(hq, (l2 + 1, vi2))
    if n % 2 == 0:
        vi = hq[0][1]
        v1 = vvv[vi].pop()
        print(f'{c + 1} {v1 + 1}')

G - Count Substring Query

• 英小文字からなる文字列 $S$ が与えられます。
• $Q$ 個のクエリが与えられるので順に処理してください。$i$ 番目のクエリは以下で説明されます。
  • 英小文字からなる文字列 $T_i$ が与えられる。
  • $S$ の部分文字列であって、$T_i$ と一致するものは何通りあるか出力する。
  • ただし、ある二つの部分文字列が文字列として同じでも、取り出された位置が異なるならばそれらは別々に数えるものとする。

• $1\leq |S|\leq 5\times 10^5$
• $1\leq Q\leq 5\times 10^5$
• $ 1\leq |T_i|\leq |S|$
• $\displaystyle \sum_{i=1}^Q |T_i|\leq 5\times 10^5$
• $S,T_i$ は英小文字列

検索文字列 $T$ が複数あるということで、Aho-Corasickが連想されるが、下手に実装するとTLEする。

■TLE解法
各 Ti の終端ノードに「自身は Ti の終端ノードである」という情報を持たせ、
S のTrie上での遷移の過程でそのノードに訪れたとき、クエリ i の答えに1加算するような実装だと、

S: aaaaaaaaaaaa...  （500000の長さ）
T: a, aa, aaa, ...  （約1000個のクエリ）

のような時、以下のようなTrie木が構築される。

(root)--(a)--(aa)--...--(a x 999)--(a x 1000),
                            ↑______________/ failure

末端の(a x 1000)は全てのTiの終端であるため、
訪れたとき、1000個のクエリの答えに加算処理が走ってしまう。
末端に到達したらループするため S の1文字毎に1000個の処理が走り、
全体で O(|S| √Σ|T_i|)、5x10^8 くらいかかってしまう。

ここでは逆に、「$S$ からの遷移の過程で、このノードは何回訪れられた」という情報 visited_count を記録する。
（下記の補正を行った上で）各 $T_i$ の終端ノードの visited_count が、そのままクエリ $i$ の答えとなる。

Aho-Corasick法では、$T_i$ の終端ノードが、他の $T_j$ 次第で $S$ に含まれているのに訪れられないことがある。

(root)----(a)----(ab)----(abc)    S = babc, T = [abc, bc, c] の場合、Sの遷移は
   |        ↙~~~'  ↙~~~'        (root)→(b)→(root)→(a)→(ab)→(abc)
   |------(b)----(bc)             の順に訪れられる
   |        ↙~~~'
   `------(c)

"abc" の方が長くマッチングしているので優先され、
"bc" "c" は含まれているが訪れられない。

これをfailure-linkを使って補正する。
$S$ からの遷移終了後、葉から「自身の failure-link 先の visited_count」に「自身の visited_count」を加算していけばよい。
上記の例の場合は、(abc) から (bc) に visited_count が加算され、 またその後 (bc) から (c) にも加算され、(bc),(c) の visited_count がちゃんと $1$ となる。
failure-linkを辿ると必ず深さが浅いノードに遷移するので、rootからの深さが深い順に行うと正しく補正できる。

計算量は $O(|S| + \sum{|T_i|})$ となる。

• maspy氏によるユーザー解説: 解説 - トヨタ自動車プログラミングコンテスト2024#7（AtCoder Beginner Contest 362）

$S$ のSuffix-Arrayを構築しておけば、ある文字列 $T$ の出現回数は、二分探索で求められる。

文字列の二分探索なので、1回毎に最大 $|T_i|$ の計算がかかり、 クエリ部分の計算量は全体で $O(\sum{|T_i|}\log{|S|})$ となるが、間に合う。

$S$ のSuffix-Array内における、$T_i$ の“bisect_left 的な位置”と、“bisect_right 的な位置”の差が答えとなる。

また、二分探索が不要な実装として、$S$ と $T$ を全て結合したSuffix-Arrayを構築する方法もある。

• MMNMM氏によるユーザー解説: 解説 - トヨタ自動車プログラミングコンテスト2024#7（AtCoder Beginner Contest 362）

各文字を $1～26$ に変換して整数リストにした上で、区切り文字として $0,27,28$ を使って、

• $U=S+(27)+T_1+(0)+T_1+(28)+T_2+(0)+T_2+(28)+...+T_Q+(0)+T_Q+(28)$
  • ※ $+$ は配列結合

のように結合した文字列 $U$ に対して Suffix-Arrayを構築する。

こうすると、必ずSuffix-Arrayの中では

• ① $(T_i,0)$ から始まるSuffix
• ② $S$ の中で $T_i$ から始まるSuffix（複数あることも、ないこともある）
• ③ $(T_i,28)$ から始まるSuffix

の順に並ぶ。①,③の $U$ 上の開始位置を求めておけばSuffix-Array上の位置がわかるため、挟まれた②の個数が答えとなる。

ただし、①と③の間には他のクエリのSuffixも含まれうるので、それは数えてはいけない。
Suffix-Array構築後、$S$ から始まるものだけを抽出して累積和などで整理しておけばよい。

計算量は $O(|S|+\sum{|T_i|}+\sigma)$（※$\sigma$: 文字の種類数=29）。

なお、ac-library for python の suffix_array() には引数に文字種数 $\sigma$ に相当する“upper”を与えられるが、これを省略すると3秒以上かかりTLEした。
ちゃんと $28$ を設定すると半分以下の1300msで通ったので、高速化のためには重要なパラメータである。

• https://atcoder.github.io/ac-library/production/document_ja/string.html

Python3

from itertools import accumulate
from atcoder.string import suffix_array


def decode(s):
    return [ord(c) - 96 for c in s]  # 1始まりとする


s = input()
n = int(input())
ttt = [input() for _ in range(n)]
BEGEN = 0
END1 = 27
END2 = 28
len_s = len(s)

s = decode(s)
s.append(END1)
begins = []
ends = []
for t in ttt:
    t = decode(t)
    begins.append(len(s))
    s.extend(t)
    s.append(BEGEN)
    ends.append(len(s))
    s.extend(t)
    s.append(END2)

res = suffix_array(s, END2)
pos = [0] * len(s)
for i, p in enumerate(res):
    pos[p] = i

counter = [0] * len(s)
for i in range(len_s):
    counter[pos[i]] += 1
counter = list(accumulate(counter))

buf = []
for i in range(n):
    b = begins[i]
    e = ends[i]
    ans = counter[pos[e]] - counter[pos[b]]
    buf.append(ans)

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