ユニークビジョンプログラミングコンテスト2022（AtCoder Beginner Contest 248）E,F,G問題メモ

ユニークビジョンプログラミングコンテスト2022（AtCoder Beginner Contest 248）

Gみたいな実装をバグらせず書きたいね。

E - K-colinear Line

• 2次元平面上に $N$ 個の点 $(x_1,y_1),(x_2,y_2),...$ がある
• $K$ 個以上の点を通る直線が何本あるか、答えよ
  • 無数にある場合は 'Infinity' と答えよ
• $1 \le N \le 300$
• 同じ座標に複数点はない

$K=1$ のとき、1点を通る直線は無限にあるのでInfinity。

以下、$K \ge 2$ とする。これならそもそも2点のペアというのが有限個で、2点決めれば直線は1つに決まるので直線も有限個。

直線の方程式は、$y=ax+b$ と表す方法がわかりやすくはあるんだけど、 垂直な線が表現できなかったり、傾き $a$ や切片 $b$ が小数になったりの問題がある。

$ax+by+c=0$ で表現すると、垂直な線も表現でき、適当に倍化することによって全パラメータを整数にできる。
この $(a,b,c)$ のパラメータ3点セットで直線が一意に決まる。

ただし、同じ直線が違うパラメータにならないよう、以下に注意する。

• $a,b$ は最大公約数が1になるようにしておく
• 正負逆転しただけのものが別々にならないよう、たとえば $a$ は必ず非負、$a=0$ のときは $b$ が正になるようにする

全2点ペアで、そこを通る直線の $(a,b,c)$ を求め、それごとに個数をカウントしていくと、 たとえば $k$ 点を通る直線なら、$\dfrac{k(k-1)}{2}$ 回カウントされるはず。

$\dfrac{K(K-1)}{2}$ 回以上カウントされている $(a,b,c)$ の個数が答え。

Python3

F - Keep Connect

• 正方形を横に並べたような形の、$2N$ 頂点 $3N-2$ 辺のグラフ（詳細は問題ページ）
• ここから $k=1～N-1$ について、以下の問題の答えを $P$ で割ったあまりを出力せよ
  • グラフからちょうど $k$ 本の辺を取り除く方法で、取り除いた後もグラフが連結であるものの個数
• $2 \le N \le 3000$
• $9 \times 10^8 \le P \le 10^9$

きちんと場合分けして遷移の考察を必要とするDP。

縦2～5 [マス or 頂点] 程度で横にだけ長い [グリッド or グラフ] 上の数え上げ問題などは、 左から順に考えて、考慮した中で最後の1列（または何列か）の取り得る状態で場合分けしたDPが解法となる場合がある。

今回の場合は、「$i$ 列目の上・下の頂点が、左上の頂点と連結か？」を考えればよさそう。

 0   1         i
●--○-- .. --◎--
｜  ｜        ｜
○--○-- .. --◎--

全体を連結にできる可能性が残っている中で、取り得る状態は以下の3つとなる。

• 0:両方連結
• 1:上だけ連結
• 2:下だけ連結

たとえば、以下のような状態なら 2:下だけ連結 となる。
だが、この状態からでも $i+1$ 番目以降で上下を連結させれば、全体を連結にすることはできる。

●--○  ○--◎
    ｜
○--○--○--◎

いずれの場合も、$i-1$ 列目以前の頂点は、左上頂点●か、$i$ 列目の頂点◎の、3つのうちいずれかには繋がっていないとダメ。
よって、$i-1$ 列目以前の状態を場合分けして管理する必要は無く、条件を満たすもののみ考えればよい。

以下のDPを考える。

• $DP[i][j=0/1/2][k]=i$ 列目までを考えて、$i$ 列目の状態が $j$ であり、削除した本数が $k$ である状態の数

初期状態は、左端の縦の1本を残すなら状態0、削除したら状態1なので、それぞれを1、他は0で初期化する。

• $DP[0][0][0]=DP[0][1][1]=1$

$i+1$ 番目への遷移は、$i→i+1$ 列目にかかる「コ」の字の3辺の削除有無を考える。全通り試しても $2^3=8$ 通り。
1個1個試していくと、連結を保ちうるのは

なので、それぞれの削除本数にも注意して遷移させていくと、 最終的に $k=1,2,...,N-1$ に対する $DP[N-1][0][k]$ のそれぞれが答えとなる。

Python3

G - GCD cost on the tree

• $N$ 頂点の木、頂点 $i$ には整数 $A_i$ が書かれている
• 2頂点 $(u,v)$ に対して、コスト $C(u,v)$ を以下で定義する
  • $F(u,v)$ を、$(u,v)$ を結ぶ単純パス上（両端含む）にある頂点数とする
  • $G(u,v)$ を、$(u,v)$ を結ぶ単純パス上の各頂点に書かれた整数 $(A_u,...,A_v)$ のGCDとする
  • $C(u,v)=F(u,v) \times G(u,v)$
• 相異なる頂点ペアは $\dfrac{N(N-1)}{2}$ 個あるが、全てについて $C(u,v)$ を足し合わせた結果を $998244353$ で割ったあまりで求めよ
• $2 \le N \le 10^5$
• $1 \le A_i \le 10^5$
• 実行制限時間: 8sec.
• メモリ制限: 2048MB（普段の2倍）

概要は公式Editorialと同じ。

実装が重ための木DP。でもわりと素直な解法ではあるので、実装力勝負、かも。

木DPで答えを求める方法を検討する。適当に根を決めて根付き木とする。

• 競技プログラミングにおける木DP問題まとめ - はまやんはまやんはまやん
• 木DPと全方位木DPを基礎から抽象化まで解説【競技プログラミング】 | アルゴリズムロジック

木DPは、各頂点 $v$ について「子から『何か』を伝えられ、他の子および $v$ 自身の情報とマージして、親に同型の『何か』を伝える」過程の、その「何か」にどういう情報を持たせ、どういう更新を行えばよいかを決めることが基本となる。

で、いろいろ考えると（雑な省略）、以下の情報があれば答えを求められることがわかる。

• $DP[v][g][j:0/1]=v$ の部分木上の頂点で、$G(u,v)$ が $g$ であるような $u$ の、
  • $j=0$: 個数
  • $j=1$: そのような全 $u$ についての $F(u,v)$ の総和

ちょっと入れ子構造がややこしいが、$g$ については取り得る値が飛び飛びなので辞書型で管理する。

グラフ例（丸付き数値はAi）

                 ㉔      {1: [3, 8],  4: [1, 2],  24: [1, 1]}
               ／  ＼
{25: [1, 1]}  ㉕    ㉘   {7: [2, 4],  28: [1, 1]}
                    /\
{21: [1, 1]}       ㉑⑦  {7: [1, 1]}

$v$ の全ての子 $u=u_1,u_2,...$ について $DP[u]$ が計算済みとして、そこから $v$ を含み、$v$ の部分木内で完結するパスのコストを求める。

大別して、以下の2種類のパスのコストを求めることになる。

• $v$ を一端とするパス
  • 「$v$ の部分木以下の頂点」と「$v$」を結ぶパス
• $v$ を経由し、$v$ の子孫同士を結ぶパス
  • 「$v$ の子 $u_1$ の部分木以下の頂点」と「$u_2(\neq u_1)$ の部分木以下の頂点」を結ぶパス

「$v$ の部分木以下の頂点」と「$v$」を結ぶパスの分のコストを求める。

説明のため、以下の補助的な情報を定義する。

• $tmp[v][u][g][j:0/1]=v$ の子のうち、$u$ について、$DP[u]$ の情報に $v$ を加えて更新したもの
  • （tmp[v][u] 以下の構造は DP[v] に同じ）

上のグラフ例で、例えば tmp[㉔][㉘] を求めるとする。

DP[㉘] = {7: [2, 4], 28: [1, 1]} なので、そこに $㉔$ を加えるなら、

• {7: [2, 4]}
  • $gcd(7,24)=1$ なので、$G(○,㉘)=7$ であったようなパスは、㉔を加えることで $G(○,㉔)=1$ になる
  • $j=0$:個数は変わらないが、$j=1$:総和については、個数分だけ、㉔を加えて1ずつ伸びる
  • {1: [2, 6]} となる
• {28: [1, 1]}
  • $gcd(28,24)=4$ なので、$G(○,㉘)=28$ であったようなパスは、㉔を加えることで $G(○,㉔)=4$ になる
  • {4: [1, 2]} となる

もし複数の更新前GCDで、更新後のGCDが同じになったら、単純に合算すればよい。

DP[㉘] の各要素について計算したものをまとめると、tmp[㉔][㉘] = {1: [2, 6], 4: [1, 2]} となる。

すると、「㉘の部分木以下と㉔とを結ぶパス」についてのコストの合計は、 tmp[㉔][㉘]のそれぞれの要素 {g: [c, t]} について、$g \times t$ を合計したものとなる。

つまり、この例では $1 \times 6 + 4 \times 2 = 14$ となる。

tmp[㉔][㉕] など㉔の他の子でも同様に求める。

「$v$ の子 $u_1$ の部分木以下の頂点」と「$u_2(\neq u_1)$ の部分木以下の頂点」を結ぶパスの分のコストを求める。

    v
   / \
 u1   u2
 /\   /\
：： ：：

更新後の $tmp[v][u_1]$ と $tmp[v][u_2]$ から求める。

両者の要素の組み合わせを二重イテレートで全部試す。
たとえば、$u_1$ 側から $\{12: [3,7]\}$ に該当する頂点、$u_2$ 側から $\{18: [5,11]\}$ に該当する頂点を結ぶとすると、

• 結んだときの $G(○,○)$ は全て $gcd(12,18)=6$

あとは、該当する全パスの $F(○,○)$ の合計値を知りたいが、定義に立ち戻るとこれはパス上の頂点数なので、

• $v$ より $u_1$ 側だけの延べ頂点数
• $v$ より $u_2$ 側だけの延べ頂点数

に分けて考えてよい。

$u_1$ 側の $3$ 個の頂点のうち任意の1つが、$u_2$ 側の $5$ 個の頂点と結ばれる5種類のパスについて、 その「$v$ より $u_2$ 側だけの延べ頂点数」は $11$ 個。（各パスに $v$ 自身を含む）
よって $3$ 個では $33$ 個。

同様に、$u_2$ 側の $5$ 個の頂点のうち任意の1つが、$u_1$ 側と結ばれる3種類のパスについて、 「$v$ より $u_1$ 側だけの延べ頂点数」は $7$ 個。$5$ 個では $35$ 個となる。

ただし、それぞれの合計には $v$ が重複している。
$3 \times 5=15$ 通りのパスに、1つずつ余分が含まれるので、これを除く。

つまり、15通りのパスの $F(○,○)$ の合計は $33+35-15=53$、コストは $6 \times 53 = 318$ となる。

一般化すると、$\{g_1: [c_1,t_1]\}$ と $\{g_2: [c_2,t_2]\}$ を結ぶパスの総コストは、以下のようになる。

• $gcd(g_1,g_2) \times (c_1t_2 + c_2t_1 - c_1c_2)$

$u_1,u_2$ についての計算が終わったら、それをマージする。 マージは単純に、同じ $g$ があったら足し合わせ、無ければ作ればよい。

tmp[v][u1]    = {1: [3, 6],  5: [2, 4],  10: [1, 1]}
tmp[v][u2]    = {1: [2, 3],  15: [1, 1]}
↓
tmp[v][u1&u2] = {1: [5, 9],  5: [2, 4],  10: [1, 1],  15: [1, 1]}

そして、tmp[v][u1&u2] と、tmp[v][u3] について同じことを繰り返していけばよい。

マージテクを使えば、多少計算量が落ちるかも。
まぁでも、8秒もあるのでその辺はサボっても間に合いそう。

ここまでの、$v$ の子孫同士を結ぶパスのコストの求め方は、 $tmp[v][u_1]$ などの代わりに $DP[u_1]$ を使っても（細かな計算は変わるが）一応可能である。

ただし、マージする部分でまずい部分が出てくる。

tmpを使えば、マージした後の要素は $A_v$ の約数に限られる（制約の中ではたかだか128個）。

一方 $DP$ をそのまま使った場合、$A_v$ の値が未反映のため、マージ後の要素数が膨れ上がる可能性がある。
このままいくつもの子を二重イテレート・マージすると、TLEとなってしまうケースが存在する。

$v$ の子孫同士を結ぶパスのコストを計算する過程で、マージを行った。

上記で述べたマージで、全ての $u_i$ の情報を1つにまとめたら、 そこに、$v$ 自身から始まるパスを表現する $\{A_v: [1, 1]\}$ をマージする。

これが $DP[v]$ となる。（$v$ が葉の場合は、$\{A_v: [1, 1]\}$ がそのまま $DP[v]$ となる）

木DPで答えを求める。頂点 $v$ について $DP[v]$ を求める処理は以下の通りである。

まず、各子 $u=u_1,u_2,...$ について、$DP[u]$ に $v$ を加えて更新した $tmp[v][u]$ を計算する。
その過程で $v$ を一端としたパスのコストが計算できる。

$tmp[v][u_1],tmp[v][u_2],...$ をマージしつつ、$v$ の子孫同士を結ぶパスのコストを計算する。

1つになるまでマージしたら、$v$ 自身から始まるパスを表現する $\{A_v: [1, 1]\}$ をマージし、$DP[v]$ とする。

木DPが最後まで行われる過程で計算されたコストの総和が、答えとなる。

Python3