AtCoder Beginner Contest 314 G問題メモ

AtCoder Beginner Contest 314

言語アップデート後のRated1発目。

ACL for Pythonなどを使えるようになり、また使える問題も出題されたものの、意外と使っている人は少なかった？
既に、より高機能 or 使い方を熟知してる自前ライブラリを持っている場合は、敢えて使う意味は薄いか。

G - Amulets

問題

$N$ 体のモンスターがいる
- $i$ 番目のモンスターは攻撃力 $A_i$ 、タイプ $B_i$
$M$ 個のお守りがあり、 $1～M$ の番号が付いている
体力 $H$ の高橋君が、お守りをいくつかを持って、体力の続く限り $i=1,2,...,N$ の順番にモンスターに対峙する
モンスター $i$ と対峙したとき、
- タイプ $B_i$ と一致する番号のお守りを持っている？
  - 持っている→攻撃を受けない
  - 持っていない→攻撃を受けて高橋君の体力が $A_i$ だけ減る
- その後、高橋君の体力が
  - 0より大きい→モンスター $i$ を倒す
  - 0以下→モンスター $i$ を倒せず終了する¹⁾
$K=0,1,...,M$ について、お守りを $K$ 種類持って出発したときに倒せるモンスターの最大数を求めよ
$1 \le M \le N \le 3 \times 10^5$
$1 \le H,A_i \le 10^9$

解法

いろいろ要素があって、何を固定すべきか迷う。

$i$ を固定して考えてみる。
タイプ（お守り番号）ごとに、それまでに出てきたモンスターの攻撃力を集計する。

   i         1      2      3      4      5
(Ai,Bi) = (10,3) (20,5) (10,1) (20,3) (20,5)

タイプ         2 4   1    3  5
攻撃力の和     0 0  10   30 40

この時、お守りを $K$ 個持っていると、攻撃力の和の、大きい方から $K$ 個分の攻撃を防ぐことができる。

逆に言うと、小さい方から累積和を取っていって、 $k$ 個目で累積和が $H$ 以上となったら、 $i$ 体目のモンスターまで倒すには $M-k+1$ 個のお守りが必要ということになる。

タイプ         2 4   1    3  5
攻撃力の和     0 0  10   30 40
累積和         0 0  10   40 80
                           ↑H=70 だったら i=5体目まで倒すには、1個以上のお守りが必要
                         ↑H=40   だったら i=5体目まで倒すには、2個以上のお守りが必要

$(i-1 体目まで倒すための必要なお守り数) ～ (i体目まで倒すために必要なお守り数-1)$ の $K$ について、答えは $i-1$ となる。
（実際は、 $k$ は $i-1$ の時と $i$ の時で異なっても1個なので、ある $i-1$ が答えとなる $K$ はたかだか1個となる（全て倒せる場合を除き））

$H$ の位置の探索と、状態（タイプ毎の攻撃力の累積和）の更新を効率的に行えるデータ構造は複数あるが、たとえば累積和の二分探索が可能なFenwickTreeを使って実装できる。

値と個数を別々に管理する必要があるため、2本の木を持つ
「攻撃力の和」としてあり得る値で座標圧縮しておく
- あり得る値は、事前に実際にシミュレーションして求める。最大 $N+1$ 個

モンスター $i$ での更新は、更新前のタイプ $B_i$ の攻撃力の和を $S$ として、 $S$ を示す位置から $S$ を引き、 $S+A_i$ を示す箇所に $S+A_i$ を足す、という2回の $O(\log{N})$ の操作により行える。（個数も同様に、-1,+1で更新できる）

(A6,B6) = (20, 1)  更新前のタイプ1の攻撃力の和S=10
                    -                            +
FWT上のindex     0  1  2  3  4          0  1  2  3  4
示す攻撃力の和   0 10 20 30 40    →    0 10 20 30 40
FWT_値           0 10  0 30 40          0  0  0 60 40
FWT_個数         2  1  0  1  1          2  0  0  2  1

その後、 $H$ で二分探索して必要お守り数を求めるのも $O(\log{N})$ なので、全部で $O(N \log{N})$ で求められる。

攻撃力の和が被るとき

2つ以上のタイプで攻撃力の和が被るとき、かつちょうどそこが $H$ の境界になった時、必要お守り数の計算に注意が必要となる。

FWT上のindex     0  1  2  3  4
示す攻撃力の和   0 10 20 30 40
FWT_値           0  0  0 60 40
FWT_個数         2  0  0  2  1
                         ↑H=40 のとき、FWT_値 上の二分探索で求まるindexは3

2つ以上のタイプが重なっている場合、その中のいくつかを倒せて、いくつかを倒せない、という事態が発生しうる。

累積和が $H$ 以上になる最初のindex（上例では $i=3$ ）を二分探索により求めたとき、

$i-1$ までのモンスターの攻撃力の和 $v$ ・個数 $c$ を求める（ $v=0,c=2$ ）
$i-1$ までのタイプを倒した後の残り体力は $H-v=40-0=40$
$i=3$ の示す攻撃力 $S_i=30$
この時、追加で倒せるモンスタータイプは $(H-v-1)/S=39/30=1$ 種類（切り捨て）
よって、この時に倒せるタイプは $c+(H-v-1)/S=2+1=3$ 種類→ $M-3=2$ 種類のお守りが必要

このように計算する必要がある。

Python3

        
          
          
              
              from operator import add, lt
 
 
class FenwickTreeInjectable:
    def __init__(self, n, identity_factory, func):
        self.size = n
        self.tree = [identity_factory() for _ in range(n + 1)]
        self.func = func
        self.idf = identity_factory
        self.depth = n.bit_length()
 
    def add(self, i, x):
        i += 1
        tree = self.tree
        func = self.func
        while i <= self.size:
            tree[i] = func(tree[i], x)
            i += i & -i
 
    def sum(self, i):
        i += 1
        s = self.idf()
        tree = self.tree
        func = self.func
        while i > 0:
            s = func(s, tree[i])
            i -= i & -i
        return s
 
    def lower_bound(self, x, lt):
        """
        累積和がx以上になる最小のindexと、その直前までの累積和
 
        :param lt: lt(a, b) で a < b ならTrueを返す関数
        """
        total = self.idf()
        pos = 0
        tree = self.tree
        func = self.func
        for i in range(self.depth, -1, -1):
            k = pos + (1 << i)
            if k > self.size:
                continue
            new_total = func(total, tree[k])
            if lt(new_total, x):
                total = new_total
                pos += 1 << i
        return pos, total
 
    def debug_print(self):
        prev = 0
        arr = []
        for i in range(self.size):
            curr = self.sum(i)
            arr.append(curr - prev)
            prev = curr
        print(arr)
 
 
n, m, h = map(int, input().split())
monsters = []
for _ in range(n):
    a, b = map(int, input().split())
    b -= 1
    monsters.append((a, b))
 
appeared_damage = {0}
type_damages = [0] * m
for a, b in monsters:
    type_damages[b] += a
    appeared_damage.add(type_damages[b])
 
damage_compress = sorted(appeared_damage)
damage_mapping = {d: i for i, d in enumerate(damage_compress)}
 
nn = len(appeared_damage)
fwt_dmg = FenwickTreeInjectable(nn, int, add)
fwt_cnt = FenwickTreeInjectable(nn, int, add)
fwt_cnt.add(0, m)
 
type_damages = [0] * m
ans = [n] * (m + 1)
j = 0
for i, (a, b) in enumerate(monsters):
    bfr_d = type_damages[b]
    bfr_i = damage_mapping[bfr_d]
    aft_d = bfr_d + a
    aft_i = damage_mapping[aft_d]
    type_damages[b] += a
    fwt_dmg.add(bfr_i, -bfr_d)
    fwt_dmg.add(aft_i, aft_d)
    fwt_cnt.add(bfr_i, -1)
    fwt_cnt.add(aft_i, 1)
    k, v = fwt_dmg.lower_bound(h, lt)
    if k == nn:
        continue
    c = fwt_cnt.sum(k - 1)
    c += (h - v - 1) // damage_compress[k]
    for t in range(j, m - c):
        ans[t] = i
    j = m - c
 
print(*ans)

            

        
      

¹⁾

ざんねん!! たかはしくんのぼうけんはここでおわってしまった!!