目次
東京海上日動 プログラミングコンテスト2020 C,D,E問題メモ
会社名の英語表記、「東京」は「Tokio」なんだね。
C - Lamps
問題
- ランプが $N$ 個、数直線上に並んでいて、ランプ $i$（$i=1～N$）は座標 $i$ にある
- ランプ $i$ が明るさ $L$ で光っているとき、座標 $i-L～i+L$ まで光が届く
- $i-L,i+L$ ちょうども範囲に含むとする
- はじめ、各ランプの明るさは $A_1,A_2,...,A_N$
- 以下の操作を $K$ 回おこなう
- 操作
- 各ランプについて、自身を含むいくつのランプから照らされているかを調べ、$B_i$ とする
- 各ランプの明るさを同時に $B_i$ に変更する
- $K$ 回の操作後のランプの明るさを求めよ
- $1 \le N,K \le 2 \times 10^5$
解法
一見、そのままやると $O(NK)$ かかりそうに見えるが、実際は回数を重ねると $\{N,N,N,...,N\}$ から動かなくなる。
$\{0,0,0,...,0\}$ から始めても結構な勢いで増えていくので、まぁ多分大丈夫でしょ、で実装する。
ちゃんと確かめるには、最も収束が遅そうなケースが $N=K=200000,A_i=0$ なので、実験すればよい。40回くらいで終わる。
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+...+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), ...$ となったとする。 また、$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)))
E - O(rand)
問題
- $N$ 個の相異なる非負整数 $A_1,A_2,...,A_N$ が与えられる
- 次の条件をともに満たすように、与えられた数の中から $1$ 個以上 $K$ 個以下の数を選ぶ方法は何通りあるか求めよ
- 条件
- 選ばれた数のビットごとの論理積（AND）は $S$
- 選ばれた数のビットごとの論理和（OR）は $T$
- $1 \le K \le N \le 50$
- $0 \le A_i,S,T \lt 2^{18}$
解法
包除原理。
事前処理
まず事前処理として、ビットを詰める。以下、「桁」は2進数での桁とする。
- $S$ のある桁に'1'があったら、もうその桁は'1'である $A_i$ しか選べない
- $T$ のある桁に'0'があったら、もうその桁は'0'である $A_i$ しか選べない
なので、条件を満たさない $A_i$ は除いておく。
さらにこれらの桁は後続の包除原理に邪魔になるので、無しにして考える。
S = 000010010
T = 011110111
xx x 1通りに決まるbit
A1= 001101101
↓
0011 1 1
↓ 圧縮
001111
事前処理後、$S$ は必ず0、$T$ は最大桁以下全てのビットがたった'1111…'のような数となる。
包除原理
（事前処理後の）$T$ を2進数で $L$ 桁とする。
問題文の条件は「$L$ 桁全てに、0も、1も、両方あるような選び方」と言い換えられる。
逆に、ある桁について見たときに、選ばれた数の全てが'0'だったり、'1'だったりするような桁がある選び方は条件を満たさない。
OK NG
10101 10101
00011 00011
01100 00101
^ ^ これらの桁が、全て0だったり、1だったりしている
条件を満たすものを直接数えるのは難しいので、全体から条件を満たさないものを引く。
すると包除原理が適用できて、以下のようになる。
答え ＝ 全体の選び方
- 少なくとも1bit、同じになってしまう選び方
+ 少なくとも2bit、同じになってしまう選び方
- 少なくとも3bit、同じになってしまう選び方
...
これはさらにbit集合に分割して、たとえば $L=3$ だったら
答え ＝ 全体の選び方
- (1桁目が同じになってしまう選び方)
- (2桁目が同じになってしまう選び方)
- (3桁目が同じになってしまう選び方)
+ (1桁目と2桁目が同じになってしまう選び方)
+ (2桁目と3桁目が同じになってしまう選び方)
+ (1桁目と3桁目が同じになってしまう選び方)
- (1桁目と2桁目と3桁目が同じになってしまう選び方)
このように、全てのbit集合についてそこが同じになってしまう選び方を数え、pop_count（bitが立っている個数）で正負を切り分ければ数え上げられる。
たとえば、$L=5$ でbit集合 01101（2桁目と3桁目と5桁目）が同じになる選び方は、各 $A_i$ と 01101 とのANDをとった結果ごとに個数を集計して、
bit = 01101
AND
A1 10101 → 00101 ┬→ 00101: 5個
A2 00111 → 00101 ┤
... ...
Ai 11100 → 01100 ┬→ 01100: 3個
... ...
5個の中から $1$ 個以上 $K$ 個以下選ぶ選び方、3個の中から $1$ 個以上 $K$ 個以下選ぶ選び方……というのを足し合わせればよい。
これは二項係数 ${}_{n}C_{r}$ を使って、$n=1～N$ ごとに前計算しておけばすぐに求まる。
そうやって足し合わせた結果が「bit集合 01101 が同じになってしまう選び方」で、今回はpop_countが3で奇数なので、答えから引き算すればよい。
計算量
- 事前処理に $O(NL)$
- 二項係数の前計算が $O(NK)$
- $L$ 桁の全てのbit集合について、各 $A_i$ とのANDを集計するところが $O(2^LN)$
全体で $O(N(K+2^L))$ でできる。
