トヨタ自動車プログラミングコンテスト2023#1(AtCoder Beginner Contest 298) G,Ex問題メモ
G - Strawberry War
問題
長方形のケーキがある。ケーキは $H$ 行 $W$ 列に並ぶ区画に分かれている
上から $i$ 行目、左から $j$ 列目の区画にはイチゴが $s_{i,j}$ 個載っている。
$T$ 回の切断を行ってケーキを $T+1$ 切れに分割する。各回の切断は、次のいずれかの方法で行う。
分割後のケーキに載ったイチゴの数をできるだけ均等にしたい。
分割後の $T+1$ 切れのケーキに載ったイチゴの個数の最大 $M$ と最小 $m$ の差 $M-m$ がとりうる最小値を求めよ。
$1 \le H,W \le 6$
$0 \le s_{i,j} \le 10^{16}$
解法
$H,W$ の制約がめちゃ小さいが、ケーキの切り方も自由度が高い。
こういうのを解くには、2次元区間DP（一般的な呼び方は知らない）が適する。
自由に切れるとはいえ、途中までナイフを入れて止める、みたいな切り方はできず、
1回の分割では必ず縦か横にナイフを通しきる必要がある。
こういうのはダメ こういうのはいい
┌─┬┐ ┌─┬┐
├┬┤│ ├┬┤│
│├┴┤ │├┤│
└┴─┘ └┴┴┘
よって、$[u,d)$ 行 $[l,r)$ 列の部分長方形は、
「全く切らない」か「2つの矩形に分割」のいずれかに限られる。
2つの長方形に分割する場合の切り口の候補は、その長方形の (高さ-1)+(幅-1) 個。
区間DP的に、小さい長方形から求めていく。
とりあえず、ちょうど $T$ 回切らないといけないので回数の情報も必要で、以下の添字を持ったDPが考えられる。
$DP[u,d,l,r,c]=[u,d)$ 行 $[l,r)$ 列の部分長方形を $c$ 回切った時に（いい感じに答えが求められる何か）
だが、この（いい感じに答えが求められる何か）の内容が意外と難しい。
2つの長方形を1つに統合した結果を、ある程度簡単に求められないといけない。
「最大と最小の差の最小化」で他によく使われる考え方として、「最小値の固定」がある。
最小値を固定すれば、その前提の中で最大値を最小化すればよくなる。
固定するべき最小値をどうやって決めればいいのか迷うが、
実は $6 \times 6$ の矩形の部分長方形は $u,d,l,r$ の決め方 $441$ 通りと意外と少なく、
最小値はこの上に載っているイチゴの数のどれかになるのは明らか。
これを全通り試せばよい。
最小値を $m$ に固定したら、先ほどと同様、全部分長方形を $u,d,l,r$ の全探索で調べて、
切らない場合 $DP[*,*,*,*,0]$ を埋められる。
この時、$m$ 未満になる場合は最小値の前提に反するので INF として、不可能であることを示しておく。
あとは、小さい矩形から順番に求めていき、$DP[0,H,0,W,T] - m$ が最小値を $m$ にした時の答え。
計算量は、最小値の固定の仕方が $O(H^2W^2)$、DPの状態数が $O(H^2W^2T)$、
遷移が「どこで切るか」と「2つの矩形の切った回数を何回と何回にするか」で $O((H+W)T)$。
$N=H=W$ としてざっくり $O(N^{13})$ となるが、
まぁ、実際は $l$ と $r$ の大小関係が決まってたり、
$l,r$ が小さい内は遷移の $T$ の部分も少ないので省略できるなど、
簡単にできる枝刈りをすれば十分通る。
Python3
import os
import sys
import numpy as np
def solve(inp):
h, w, t = inp[:3]
sss = np.zeros((h, w), dtype=np.int64)
for i in range(h):
sss[i] = inp[3 + i * w:3 + (i + 1) * w]
INF = 1 << 60
min_cand = set()
for u in range(h):
for d in range(u + 1, h + 1):
for l in range(w):
for r in range(l + 1, w + 1):
area = sss[u:d, l:r].sum()
min_cand.add(area)
ans = INF
for mn in min_cand:
dp = np.full((h, h + 1, w, w + 1, t + 1), INF, dtype=np.int64)
for u in range(h):
for d in range(u + 1, h + 1):
for l in range(w):
for r in range(l + 1, w + 1):
area = sss[u:d, l:r].sum()
if area >= mn:
dp[u, d, l, r, 0] = area
for height in range(1, h + 1):
for width in range(1, w + 1):
for u in range(h - height + 1):
d = u + height
for l in range(w - width + 1):
r = l + width
for k1 in range(t):
for k2 in range(t - k1):
for m in range(u + 1, d):
dp[u, d, l, r, k1 + k2 + 1] = min(dp[u, d, l, r, k1 + k2 + 1],
max(dp[u, m, l, r, k1],
dp[m, d, l, r, k2]))
for m in range(l + 1, r):
dp[u, d, l, r, k1 + k2 + 1] = min(dp[u, d, l, r, k1 + k2 + 1],
max(dp[u, d, l, m, k1],
dp[u, d, m, r, k2]))
ans = min(ans, dp[0, h, 0, w, t] - mn)
return ans
SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
from numba.pycc import CC
cc = CC('my_module')
cc.export('solve', SIGNATURE)(solve)
cc.compile()
exit()
if os.name == 'posix':
# noinspection PyUnresolvedReferences
from my_module import solve
else:
from numba import njit
solve = njit(SIGNATURE, cache=True)(solve)
print('compiled', file=sys.stderr)
inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)
解法2
2次元区間DPは同じだが、「最大値」と「最小値」を両方固定する場合も考えられる。その上で、
$DP[u,d,l,r]=[u,d)$ 行 $[l,r)$ 列の長方形を、全ての区画を最小～最大の範囲に収めつつ、何回切断できるかの組
をbitsetで持っておく等でも解けるみたい。
計算量は多少増えるため、最大や最小の時点で明らかに不可能なものはスキップするなどの枝刈りが必要となる。
Ex - Sum of Min of Length
問題
$N$ 頂点の木が与えられる
頂点 $u,v$ 間の距離（パスの辺数）を $d(u,v)$ で表すとする
$Q$ 個のクエリに答えよ
$L,R$ が与えられるので、$\displaystyle \sum_{v=1}^N \min(d(v,L),d(v,R))$ を出力する
つまり、全ての頂点についての、$L$ か $R$ 近い方との距離の総和
$1 \le N,Q \le 2 \times 10^5$
解法
「$R$ の方が近い頂点」「$L$ の方が近い頂点」は、$L$ と $R$ の真ん中の点または辺が境界となってはっきり分かれる。
●-❶-, ↓ ,-○ （フォントの都合上、Lを❶、Rを②と表現）
●-●--◎--○-②-○
●-●-' / `-○
◎-◎
真ん中が点の場合、その点、およびその点から出てる他の枝（◎）はどっちでも一緒なので、どっちかに属させればよい。
（ここでは以降、●に属させる）
で、❶から各●への距離の和と、②から各○への距離の和を別々に求めて合わせれば答えとなる。
❶か②を根としたとき、自分に属さない方の頂点は、中央（付近）の頂点を根とする部分木の形になる。
❶
/\
●●
/｜\
●●❹←中央
｜｜＼
●● ③ ③以下の部分木は②の方に属する
｜ /\
●○②
/\
○○
なので、
が高速に求められれば、答えの半分を出せる。②も同様に求めて合わせれば答え（その場合、引くべき部分木の根は❹）。
❶を根として、各頂点 $v$ の以下の情報が計算されているとする。
自身から、自身を根とする部分木の各頂点との距離の総和 $ds_{1,v}$
自身を根とする部分木の頂点数 $cnt_{1,v}$
根から自身までの距離 $d_{1,v}$
すると、以下が成り立つ。
（❶から③を根とする部分木との距離の総和）＝（③から③を根とする部分木との距離の総和 $ds_{1,3}$）+（③を根とする部分木の頂点数 $cnt_{1,3}$）×（❶と③の距離 $d_{1,3}$）
❶から他の全頂点への距離は $ds_{1,1}$ なので、これで答えの半分が出せた。
今は❶を根としたが、$ds,cnt$ はクエリの各 $L,R$ を根とした場合について必要で、これを個別に求めていれば間に合わない。
だが、実際は、根を1つ隣の頂点に移したとき、変わる部分というのは限られる。（更新の仕方は省略）
❶
/＼
❺ ●
/\ ：
●●
：：
$ds_{5,v}$ は、$ds_{5,1}$ と $ds_{5,5}$ を除き、$ds_{1,v}$ と変わらない
$cnt_{5,v}$ も、$cnt_{5,1}$ と $cnt_{5,5}$ を除き、$cnt_{1,v}$ と変わらない
なので、まず根❶について求めた後、オイラーツアーの順番で根を移動させながら $ds, cnt$ の一部だけを更新する、
ということを繰り返せば、全頂点を根としたときの $ds, cnt$ が求められる。
$d_{u,v}$ は、根を更新すればごそっと1増えたり減ったりで変わってしまうのでこの更新には向かないが、
これは各クエリにつき中央の点を求めるときに一緒に求めておけばよい。
まとめ
木からダブリングを構築し、LCAや“$k$ 遡った頂点” を高速に求められるようにしておく
クエリ $L,R$ につき、以下の事前加工をおこなう
LCAを求め、距離を求め、中央の点（$L,R$ それぞれを根としたときに除く部分木の根）を求める
各頂点に付き、（何番目のクエリで必要か, 除く部分木の根 $w$, 自身と $w$ との距離）を登録する
1を根として $ds_{1,v},cnt_{1,v}$ を求める
オイラーツアー順に根を移動させながら、クエリで必要な頂点が根になったタイミングで、事前加工で登録した情報を元に、答えの半分を計算する
計算量は、何回も木を走査するので定数倍はそれなりにかかるが、
ダブリングの部分で $O(N\log{N})$、木DPで $O(N)$、オイラーツアーで $O(N)$ となる。
Python3
class LcaDoubling:
"""
links[v] = { u1, u2, ... } (u:隣接頂点, 親は含まない)
というグラフ情報から、ダブリングによるLCAを構築。
任意の2頂点のLCAを取得できるようにする
"""
def __init__(self, n, links, root=0):
self.depths = [-1] * n
prev_ancestors = self._init_dfs(n, links, root)
self.ancestors = [prev_ancestors]
max_depth = max(self.depths)
d = 1
while d < max_depth:
next_ancestors = [prev_ancestors[p] for p in prev_ancestors]
self.ancestors.append(next_ancestors)
d <<= 1
prev_ancestors = next_ancestors
def _init_dfs(self, n, links, root):
q = [root]
direct_ancestors = [-1] * (n + 1) # 頂点数より1個長くし、存在しないことを-1で表す。末尾(-1)要素は常に-1
self.depths[root] = 0
while q:
u = q.pop()
for v in links[u]:
if v == direct_ancestors[u]:
continue
direct_ancestors[v] = u
self.depths[v] = self.depths[u] + 1
q.append(v)
return direct_ancestors
def get_lca(self, u, v):
du, dv = self.depths[u], self.depths[v]
if du > dv:
u, v = v, u
du, dv = dv, du
tu = u
tv = self.upstream(v, dv - du)
if u == tv:
return u
for k in range(du.bit_length() - 1, -1, -1):
mu = self.ancestors[k][tu]
mv = self.ancestors[k][tv]
if mu != mv:
tu = mu
tv = mv
lca = self.ancestors[0][tu]
assert lca == self.ancestors[0][tv]
return lca
def upstream(self, v, k):
i = 0
while k:
if k & 1:
v = self.ancestors[i][v]
k >>= 1
i += 1
return v
def calc_dist_sum(n, links, root):
subtree_counts = [0] * n
subtree_dist_sums = [0] * n
status = [0] * n
q = [root]
while q:
u = q[-1]
if status[u] == 0:
status[u] = 1
for v in links[u]:
if status[v] == 0:
q.append(v)
elif status[u] == 1:
q.pop()
status[u] = 2
vc = 1
ds = 0
for v in links[u]:
if status[v] == 2:
vcv = subtree_counts[v]
dsv = subtree_dist_sums[v]
vc += vcv
ds += dsv + vcv
subtree_counts[u] = vc
subtree_dist_sums[u] = ds
return subtree_counts, subtree_dist_sums
n = int(input())
links = [[] for _ in range(n)]
for _ in range(n - 1):
u, v = map(int, input().split())
u -= 1
v -= 1
links[u].append(v)
links[v].append(u)
lca = LcaDoubling(n, links, 0)
parents = lca.ancestors[0]
qn = int(input())
queries = [[] for _ in range(n)]
for i in range(qn):
u, v = map(int, input().split())
u -= 1
v -= 1
if u == v:
queries[u].append((i, -1, -1))
continue
m = lca.get_lca(u, v)
du = lca.depths[u]
dv = lca.depths[v]
if du > dv:
u, v = v, u
du, dv = dv, du
path_length = du + dv - 2 * lca.depths[m]
hpl = (path_length - 1) // 2
m1 = lca.upstream(v, hpl)
m2 = parents[m1]
queries[u].append((i, m1, path_length - hpl))
queries[v].append((i, m2, hpl + 1))
subtree_counts, subtree_dist_sums = calc_dist_sum(n, links, 0)
current_root = 0
progress = [0] * n
limits = list(map(len, links))
ans = [0] * qn
def reroot(u, v):
# 根uの子であるvを新たな根とする
scv = subtree_counts[v]
dsv = subtree_dist_sums[v]
subtree_counts[u] -= scv
subtree_counts[v] = n
subtree_dist_sums[u] -= dsv + scv
subtree_dist_sums[v] += subtree_dist_sums[u] + subtree_counts[u]
while current_root != 0 or progress[0] < limits[0]:
u = current_root
if progress[u] == 0:
for qi, m, d in queries[u]:
if m == -1:
ans[qi] += subtree_dist_sums[u]
else:
ans[qi] += subtree_dist_sums[u] - subtree_dist_sums[m] - subtree_counts[m] * d
if progress[u] >= limits[u]:
if parents[u] != -1:
reroot(u, parents[u])
current_root = parents[u]
continue
v = links[u][progress[u]]
progress[u] += 1
if v == parents[u]:
continue
# vを新たな根とする
reroot(u, v)
current_root = v
print('\n'.join(map(str, ans)))
