ユニークビジョンプログラミングコンテスト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

from collections import defaultdict
from math import gcd

n, k = map(int, input().split())

if k == 1:
    print('Infinity')
    exit()

points = [tuple(map(int, input().split())) for _ in range(n)]
points.sort()
lines = defaultdict(int)
for i in range(n):
    x1, y1 = points[i]
    for j in range(i + 1, n):
        x2, y2 = points[j]
        a = y2 - y1
        b = x1 - x2
        g = gcd(a, b)
        a //= g
        b //= g
        c = a * x1 + b * y1
        lines[a, b, c] += 1

thr = k * (k - 1) // 2
ans = sum(cnt >= thr for cnt in lines.values())
print(ans)

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

import numpy as np

n, p = map(int, input().split())
# 0: 両方連結
# 1: 上は左上端と連結、下はこれ以降のどこかで連結させないといけない
# 2: 下は左上端と連結、上はこれ以降のどこかで連結させないといけない

dp = np.zeros((3, n), np.int64)
dp[0, 0] = 1
dp[1, 1] = 1

for _ in range(n - 1):
    ndp = np.zeros_like(dp)
    ndp[0] = (dp[0] + dp[1] + dp[2]) % p
    ndp[0, 1:] += dp[0, :-1] * 3
    ndp[1, 1:] += dp[1, :-1]
    ndp[2, 1:] += dp[2, :-1]
    ndp[1, 2:] += dp[0, :-2]
    ndp[2, 2:] += dp[0, :-2]
    dp = ndp % p

print(*dp[0, 1:])

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

from heapq import heapify, heappop, heapreplace
from math import gcd


def merge(counts, a, total, count):
    if a in counts:
        counts[a][0] += total
        counts[a][1] += count
    else:
        counts[a] = [total, count]


n = int(input())
aaa = list(map(int, input().split()))
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)

count = [{} for _ in range(n)]
MOD = 998244353
status = [0] * n
ans = 0

q = [0]
while q:
    u = q[-1]
    if status[u] == 0:
        status[u] = 1
        p = -1
        for v in links[u]:
            if status[v] == 0:
                q.append(v)
            else:
                p = v
        links[u].discard(p)
    else:
        q.pop()
        a = aaa[u]
        children = []
        hq = []
        for v in links[u]:
            ch = {}
            for b, item in count[v].items():
                g = gcd(a, b)
                s = item[0] + item[1]
                ans += g * s
                ans %= MOD
                merge(ch, g, s, item[1])
            hq.append((len(ch), len(children)))
            children.append(ch)

        if len(children) > 1:
            heapify(hq)
            while len(hq) > 1:
                l1, ci1 = heappop(hq)
                l2, ci2 = hq[0]
                ch1 = children[ci1]
                ch2 = children[ci2]

                for b, item1 in ch1.items():
                    for c, item2 in ch2.items():
                        g = gcd(b, c)
                        ans += g * (item1[0] * item2[1] + item1[1] * item2[0] - item1[1] * item2[1])
                        ans %= MOD

                for b, item1 in ch1.items():
                    merge(ch2, b, *item1)

                heapreplace(hq, (len(ch2), ci2))

            ch = children[hq[0][1]]
            merge(ch, a, 1, 1)
            count[u] = ch

        elif len(children) == 1:
            ch = children[0]
            merge(ch, a, 1, 1)
            count[u] = ch
        else:
            count[u] = {a: [1, 1]}

print(ans)