−目次

東京海上日動プログラミングコンテスト2020 C,D,E問題メモ
- C - Lamps
  - 問題
  - 解法
- D - Knapsack Queries on a tree
  - 問題
  - 解法
- E - O(rand)
  - 問題
  - 解法

東京海上日動プログラミングコンテスト2020 C,D,E問題メモ

東京海上日動プログラミングコンテスト2020

会社名の英語表記、「東京」は「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))$ でできる。

実装

事前処理で実際にbitを詰めるのは面倒なので、イテレートするbit集合を限定する方が楽。

詰めた後に残るbitというのは、詰める前の「 $S$ と $T$ のXOR」なので、その部分集合を列挙するとよい。

（厳密には $S=1$ かつ $T=0$ のbitがあった場合に異なるが、この場合、答えは0で、最初にフィルタリングしたときに全ての $A_i$ は消えるので、問題にならない）

ビット演算: bitで表された部分集合の列挙の仕方

Python3

−目次

東京海上日動 プログラミングコンテスト2020 C,D,E問題メモ