東京海上日動 プログラミングコンテスト2020 C,D問題メモ
会社名の英語表記、「東京」は「Tokio」なんだね。
C - Lamps
問題
- ランプが $N$ 個、数直線上に並んでいて、ランプ $i$（$i=1～N$）は座標 $i$ にある
- ランプ $i$ が明るさ $L$ で光っているとき、座標 $i-L～i+L$ まで光が届く
- $i-L,i+L$ ちょうども範囲に含むとする
- はじめ、各ランプの明るさは $A_1,A_2,\ldots,A_N$
- 以下の操作を $K$ 回おこなう
- 操作
- 各ランプについて、自身を含むいくつのランプから照らされているかを調べ、$B_i$ とする
- 各ランプの明るさを同時に $B_i$ に変更する
- $K$ 回の操作後のランプの明るさを求めよ
- $1 \le N,K \le 2 \times 10^5$
解法
一見、そのままやると $O(NK)$ かかりそうに見えるが、実際は回数を重ねると $\{N,N,N,\ldots,N\}$ から動かなくなる。
$\{0,0,0,\ldots,0\}$ から始めても結構な勢いで増えていくので、まぁ多分大丈夫でしょ、で実装する。
ちゃんと確かめるには、最も収束が遅いケースがわかりやすく $N=K=200000,A_i=0$ なので、そのテストケースを作って実行すればよい。
import sys
from itertools import accumulate
n, k, *aaa = map(int, sys.stdin.buffer.read().split())
for _ in range(k):
bbb = [0] * (n + 1)
for i in range(n):
a = aaa[i]
bbb[max(0, i - a)] += 1
bbb[min(n, i + a + 1)] -= 1
aaa = list(accumulate(bbb))
aaa.pop()
if all(a == n for a in aaa):
break
print(*aaa)
D - Knapsack Queries on a tree
問題
- 頂点数 $N$ の二分木がある
- 根は頂点1、頂点 $i$ の子は（存在すれば）$2i,2i+1$
- 各頂点 $i$ には、重さ $W_i$、価値 $V_i$ が定められている
- クエリが $Q$ 個ある。それぞれに答えよ
- クエリ
- $j$ 番目のクエリは、頂点 $U_j$ と容量 $L_j$ からなる
- $U_i$ から根までのパスに含まれる頂点の集合で、容量 $L_j$ の0-1ナップサック問題を解け
- $1 \le N \lt 2^{18}$
- $1 \le Q \le 10^5$
- $1 \le W_i,V_i,L_j \le 10^5$
解法
2種類のナップサックDPと半分全列挙。
雑な愚直解の計算量見積もり
ナップサック問題は一般に、「$dp[v]=$ ある価値 $v$ を達成できる最小重さ」または「$dp[w]=$ ある重さ $w$ で作れる最大価値」として埋めていくため、 価値か重さ、どちらかをキーとすることになる。
今回は重さをキーとする。上限は $L_{max}=10^5$ である。
二分木の高さは最大で $D=18$ なので、1回のクエリでは最大18回の更新が行われる。
クエリ $Q$ 回で $O(QDL_{max})$ となり、最大値の代入で $1.8 \times 10^{11}$ かかる。余裕で間に合わない。
メモ化
根である頂点1は毎回候補に入る。また、その下の頂点2,3も、まぁほぼ確実に各クエリでどちらかは候補に入る。
根からDPを更新することにすれば、ある頂点まで計算した結果は複数のクエリで使い回せる。これを保存しておけばよさそう。
ただ、頂点数が $2^{18}=262144$ でクエリが $10^5$ 個なので、クエリの配置をばらけさせると、ほぼ全ての頂点が1度は使われるようなテストケースにすることも可能になってしまう。
仮に頂点 $2^{17}$ 個のそれぞれに長さ $10^5$ の配列を記録していたら、intを4byteとしてもそれだけで52GB、MLEが発生する。
従って、ある程度の深さ以降はメモ化を打ち切り、複数回参照されることがあっても毎回計算する、といった工夫が必要となる。
DPの分割
さて、ある深さ以降でDPの記録を止める場合、そこでDPの計算方法も分けると、計算量の削減になる。
上側のDP、下側のDPと呼ぶことにする。
一般に、DPの実装には、配列で持つものと、辞書で持つものの2種類ある。 この問題は、2つのメリット・デメリットを上手く使い分ける問題でもある。
- 配列
- 長さ $L_{max}$ の配列、はじめ0で初期化
- $w,v$ で更新するとき、$i = 0～L_{max}-w$ につき、$DP[i+w] ← DP[i]+v$
- $←$ は、最大を更新する操作
- 計算量は1アイテム毎に常に $L_{max}$
- 最終的にぴったりでなくても、$DP[W]$ を参照すれば $W$ 以下の最適解が入っている
- 辞書
- はじめ $\{0: 0\}$ で初期化
- $w,v$ で更新するとき、辞書内の各要素 $cw, cv$ につき、$DP[cw+w] ← cv+v$
- $←$ はキーがなければ作成、既にあれば最大を更新する操作
- 計算量は1から始まって1アイテム毎に最悪倍々になり、最大で $L_{max}$
- 重さがぴったり $W$ になる場合の価値しかわからないので、$W$ 以下の最適解を得るには全要素のminを取る操作が必要
様々な $L_i$ に対して高速に答えを返さなければならない上側のDPには、配列を用いた方がよい。 計算量はかかるが、どこを参照しても最適解が入っている、というのはメモ化に適している。
一方、下側のDPはクエリ毎に1回、合計 $Q$ 回行われるので、1回に $L_{max}=10^5$ かけてたらTLEになってしまう。
辞書を用いると、アイテム数が小さい内は計算量が少なく済む。
たとえば深さ9で分けた場合、下側のDPの更新回数は最大 $18-9=9$ 回となり、計算量は約 $2^1+2^2+\ldots+2^9=1023$ となる。 これを配列で実装していたら、$9 \times 10^5$ 回となるので、その差は大きい。
2つのDPから答えを求める
頂点 $U_j$ が上側のDPに属するなら、メモ化されているので $DP_{U_j}[L_j]$ を参照すればよい。
下側に属する場合、$U_i$ から根に向かって上下の境界まで辞書型のDPを行う。
その結果、取り得る (重さ合計, 価値合計) のパターンが $(w_1, v_1), (w_2, v_2), \ldots$ となったとする。 また、$U_j$ の祖先で、上側のDPに属する最も深い頂点を $U'_j$ とする。
ある $(w_i, v_i)$ に対して、上側で使える残り重量は $L_j-w_i$ となるので、上下合わせて $v_i+DP_{U'_j}[L_j-w_i]$ が最適解となる。
各 $(w_i,v_i)$ についての最適解を求め、その最大値が答えとなる。
計算量
入力の受け取りに $O(N)$ かかる。
二分木の最大深さを $D$、上下でDPを分割する深さを $d$ とする。
上側のDPの頂点は $2^d-1$ 個あり、クエリ次第で全て計算される可能性があるので $O(2^dL_{max})$ かかる。これは全てのクエリを通じて1回計算すればよい。
下側のDPは1回のクエリにつき、辞書型DPを用いることで $2^{D-d+1}$ の計算で $(w_i, v_i)$ を列挙できる。 上側のDPが事前計算済みという前提で、それぞれの残り重量から最適値を得るのは $O(1)$ でできるので、$O(2^{D-d})$ となる。
これらをあわせて $O(N+2^dL_{max}+2^{D-d}Q)$、$d=9$ として最大値を入れると $1.024 \times 10^8$、制限時間が3secなので何とか現実的になる。
アドホックな改善
それでも計算量的に厳しくて、PyPyを用いても、細かな高速化が必要となる。
下側のDPを辞書で行うといったが、アイテム数が9個程度なので、 実際には、各アイテムを選ぶ・選ばないの組み合わせを全て試した $2^k$ 通りの $(w_i, v_i)$ の集合（合計重さが被っても気にしない）とした方が、 キーが辞書にあるかなどの処理が省略され単純になる分、却って速くなる。
また、DPを分割する深さは、丁度真ん中の9より、10にした方が、今回のテストケースでは高速となった。11にしたらMLEだった。
import sys
def get_solve():
cache = [None] * precalc_limit
cache[0] = [0] * (weight_limit + 1)
_vvv = vvv
_www = www
_weight_limit = weight_limit
def _f(u):
""" u < precalc_limit """
if cache[u] is not None:
return cache[u]
dp = _f(u >> 1).copy()
v = _vvv[u]
w = _www[u]
for x in range(_weight_limit, w - 1, -1):
nv = dp[x - w] + v
if dp[x] < nv:
dp[x] = nv
cache[u] = dp
return dp
return _f
n, *inp = map(int, sys.stdin.buffer.read().split())
vvv = [0] + inp[0:n * 2:2]
www = [0] + inp[1:n * 2:2]
weight_limit = 10 ** 5
precalc_limit = min(1 << 10, n + 1)
solve = get_solve()
buf = []
mp = iter(inp[n * 2 + 1:])
for u, l in zip(mp, mp):
if u < precalc_limit:
dp = solve(u)
buf.append(dp[l])
continue
dp_w = [0]
dp_v = [0]
while u >= precalc_limit:
v = vvv[u]
w = www[u]
for i in range(len(dp_w)):
nw = dp_w[i] + w
if nw > l:
continue
nv = dp_v[i] + v
dp_w.append(nw)
dp_v.append(nv)
u >>= 1
ans = 0
dp = solve(u)
for w, v in zip(dp_w, dp_v):
nv = v + dp[l - w]
if ans < nv:
ans = nv
buf.append(ans)
print('\n'.join(map(str, buf)))
