NS Solutions Corporation Programming Contest 2023（AtCoder Beginner Contest 303）F,G,Ex問題メモ

NS Solutions Corporation Programming Contest 2023（AtCoder Beginner Contest 303）

F - Damage over Time

問題

$N$ 種類の魔法で、体力 $H$ のモンスターに攻撃して0以下にしたい
1ターンに1個の魔法を選んで使える
魔法 $i$ をターン $t$ に使用すると、$t$ から $t+T_i-1$ の $T_i$ ターンにかけて、1ターン毎に $D_i$ のダメージを与える
- 効果は加算で累積して、次ターンに魔法 $j$ を使うと $D_i+D_j$ のダメージを与える（魔法 $i$ が切れていなければ）
最も早く倒すには何ターンかかるか、求めよ
$1 \le N \le 3 \times 10^5$
$1 \le H \le 10^{18}$
$1 \le T_i,D_i \le 10^9$

解法

二分探索が浮かんだが、無くても解けるみたい。でもせっかくだし二分探索解法をメモ。

魔法 $i$ は、時間いっぱいかければ計 $T_iD_i$ のダメージを与えるが、、、

H = 100

10000 ターン毎ターン   1 のスリップダメージ
    1 ターン毎ターン 100 のスリップダメージ

T*D が小さくても、2番目を使った方がいいのは明らか

ゴールが見えていないと、どれくらい先まで有効な魔法を使っていいのか、見えづらい。

なので、ゴールを固定し「$K$ ターンまでに倒せるか？」の判定問題が $O(N)$ くらいで解ければ、二分探索で $O(N\log{H})$ で求められる。

判定問題

いま、$K$ まで残り $r$ ターンとする。

$T_i$ が $r$ 以上か、未満かで分ける。（その魔法の集合をそれぞれ $X,Y$ とする）

$X$ から選んで使うとしたら、有効なターン数は $r$ なので、計 $rD_i$ の総ダメージ
- → $D_i$ 最大のものを使うべき
$Y$ から選んで使うとしたら、フル（$T_iD_i$ の総ダメージ）で与えられる
- → $T_iD_i$ 最大のものを使うべき

総ダメージの大きい方を使う。図化すると、

D ↑
  ｜───┐
  ｜      │
  ｜───┼──┐
  ｜───┼──┼──┐
  ＋----------------|-----> T,残りターン
                    r

各長方形が魔法で、r より左のやつは面積が一番大きいの、
                    右にはみ出てるやつは一番高いのを選んで使う

魔法を1回使うと、$r$ が1減る。同時に $X$ の魔法で与えられる総ダメージも減る。

$H$ が馬鹿でかいので、1ターンずつこれを処理しているとTLE。
変曲点毎にまとめて処理する。（変曲点：使う魔法を変えた方が得になる可能性がある点）

変曲点は、以下の2つ。

はじめ、$X$ を使った方がダメージが大きかったが、途中から $Y$ を使った方がよくなる $r$
$Y$ から $X$ に移る魔法がある $r$

1つめは、$T_yD_y/D_x$ で出せる。
2つめも、あらかじめソートしておけば次に来る $T_i$ は分かる。

このような変曲点は、ざっくり見積もりで最大 $2N$ 個あるので、計算オーダーもその程度となる。

細かな注意点として、

$T_iD_i$ 基準と $T_i$ 基準でそれぞれソートした配列がほしくなるが、これは二分探索で毎回ソートしていると $O(\log{N})$ が余分にかかってしまう。結果は不変なので、全体で1回ソートしたものを使い回せる
$X$ の魔法を繰り返し使う場合、有効ターンが1ターンずつ減っていくので、延べターン数は等差数列の和になる

Python3

なんか、PyPyでずっとTLEが取れなくて、Numbaにしたら問題なく通った。
PyPy、たまに意図せず妙に時間がかかる地雷を踏むことがある。

import os
import sys

import numpy as np


def solve(inp):
    n, h = inp[:2]
    ttt = inp[2::2]
    ddd = inp[3::2]

    magics_by_full_damage = []
    magics_by_turn_damage = []
    for i in range(n):
        magics_by_full_damage.append((ttt[i] * ddd[i], -ttt[i], ddd[i]))
        magics_by_turn_damage.append((ddd[i], -ttt[i]))

    magics_by_full_damage.sort()
    magics_by_turn_damage.sort()

    def calc1(a, b, d, h):
        if a == 0 or b == 0 or d == 0:
            return h
        win = (h * 2 - 1) // d + 1
        if a > win // b:
            return -1
        return h - a * b // 2 * d

    def calc2(a, b, h):
        if a == 0 or b == 0:
            return h
        if a > h // b:
            return -1
        return h - a * b

    def check(h, m):
        magics_full = []
        max_part_d = 0
        for i in range(n):
            m_item3 = magics_by_full_damage[i]
            if -m_item3[1] >= m:  # t
                max_part_d = max(max_part_d, m_item3[2])  # d
            else:
                magics_full.append(m_item3)

        magics_turn = []
        for i in range(n):
            m_item2 = magics_by_turn_damage[i]
            if -m_item2[1] < m:  # t
                magics_turn.append(m_item2)

        fi = len(magics_full) - 1
        ti = len(magics_turn) - 1

        while m > 0:

            while ti >= 0 and m <= -magics_turn[ti][1]:
                max_part_d = max(max_part_d, magics_turn[ti][0])
                ti -= 1

            while fi >= 0 and (m <= -magics_full[fi][1] or magics_full[fi][2] <= max_part_d):
                max_part_d = max(max_part_d, magics_full[fi][2])
                fi -= 1

            if ti >= 0:
                _, tt = magics_turn[ti]
                tt = -tt
            else:
                return calc1(m + 1, m, max_part_d, h) <= 0

            if fi >= 0:
                ftd, ft, _ = magics_full[fi]
                ft = -ft
            else:
                return calc1(m + 1, m, max_part_d, h) <= 0

            if calc1(m + 1, m, max_part_d, h) <= 0:
                return True

            if calc2(ftd, m - ft + 1, h) <= 0:
                return True

            t = max(ft, tt)

            # print(m, h, t, max_part_d)

            if max_part_d > 0:
                sw = ftd // max_part_d
            else:
                sw = m

            if sw <= t:
                k = m - t
                h = calc1(m + m - k + 1, k, max_part_d, h)
            elif sw < m:
                k = m - sw
                h = calc1(m + m - k + 1, k, max_part_d, h)
                h = calc2(ftd, sw - t, h)
            else:
                h = calc2(ftd, m - t, h)

            if h <= 0:
                return True

            m = t

        return False

    l = 0
    r = h
    while l + 1 < r:
        m = (l + r) >> 1

        if check(h, m):
            r = m
        else:
            l = m

    return r


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

解法1と近い考え方で、二分探索が不要であるようにできる。（むしろ二分探索考察時に、何故これに気がつかないのか）

ターンを逆順から考える。
先ほどの説明中の、残りターン数 $r$ を、$r=1$ から徐々に増やしていく感じ。

「今が何ターン目かは分からないが、倒しきる最終ターン」であると考える（ちょっと奇妙だが）。
$r=1,2,3,...$ の時々に最適な魔法は決まるので、その総ダメージの合計値が $H$ になった時が、倒せるターン。

1魔法当たりの総ダメージのグラフは以下のような感じになる。

f(i,r): 残りターン r で放った魔法 i が与える総ダメージ

 総ダメ↑
       ｜
 Dk*Tk_｜              ____↙魔法 k の線
 Dj*Tj_｜       _____-~____↙魔法 j の線    斜線の傾きが各 $D$ を表す。
 Di*Ti_｜  ___／__-~_______↙魔法 i の線
       ｜ / ／_-~
       ｜/／-~
       ┼-------------------> r
          Ti   Tj     Tk

$r$ を増やしてって、$r$ ごとに $\max_i f(i,r)$ を合計していき、合計が $H$ になればその時の $r$ が答え。

解法1と同様、変曲点（グラフの屈折点、または最適な魔法が変わる交点）毎にまとめれば、 $O(N)$ で $r$ までの総ダメージを計算できる。

G - Bags Game

問題

$N$ 個のバッグが横一列に並び、$i$ 番目には $X_i$ 円入っている
太郎と次郎が、太郎を先手として交互に、以下の3つから1つ選んで行動する
- 残っている両端のどちらかのバッグから1個取る
- $A$ 円失う。「残っている両端のどちらかのバッグから1個取る」ことを $B$ 回まで繰り返す
- $C$ 円失う。「残っている両端のどちらかのバッグから1個取る」ことを $D$ 回まで繰り返す
バッグがなくなるまで繰り返し、最終的に「バッグから獲得した金額 - 失った金額」を利得とする
- （※利得が負になるだけで、2,3番目の行動はいつでも取れる）
「自分の利得 - 相手の利得」を「スコア」とし、両者、これを最大化するように行動する
太郎が実現できるスコアの最大値を求めよ
$1 \le N \le 3000$
$1 \le X_i,A,C \le 10^9$

解法

バッグは常に連続した区間が残り続ける。問題設定的にも、制約的にも、いかにも区間DPっぽい。

行動指針や取れる行動が、太郎も次郎も変わらないので、「ある盤面から、次に行動する方（先手）が実現できるスコア」は、太郎も次郎も一緒。

また、行動指針が「自分 - 相手」なので、
「ある行動をして、その後、ある盤面を相手に引き渡したときのスコア」は、
「その行動によって得られた利得 - その盤面から先手が実現できるスコア」となる。

以下のDPを考える。

$DP[w,l]=$ 長さ $w$、左端 $l$ の区間のバッグから始めて、先手が実現できるスコア

（通常、区間DPは $[l,r)$ の区間を添え字 $(l,r)$ で表現することもあるが、今回はこっちの方がやりやすい）
$r=l+w$ となる。

初期値を

$w=0$ の場合は $DP[0,*]=0$
$w=1$ の場合は $DP[1,l]=X_l$

として、狭い区間から、より広い区間を求めていく。

$[l,r)$ の時に取れる行動は、

      l                    r
...   o  o  o  o  o  o  o  ...

      x  o  o  o  o  o  o    左端を取る
      
      o  o  o  o  o  o  x    右端を取る
      
      x  x  o  o  o  x  x    A円払ってB個取る(この例では4)
      
      x  o  o  x  x  x  x    C円払ってD個取る(この例では5)

それぞれ、「x のバッグに入っている金額の和 - 支払った金額 - 残った区間のDP結果」がスコアになるので、その中の最大値を取れば良い。

ただ、$A$ 円払う行動では、どの区間を残すかが $B+1$ 通りある。

制約的に、$O(N^2)～O(N^2 \log{N})$ くらいまでしか許容されないので、これを1つ1つ調べることは $O(N^3)$ となりできない。
これを効率的に計算したい。

高速化

$DP[l,r]$ を求めるに当たっての行動 $(A,B)$ の最適なスコア計算を考える。（$(C,D)$ も同様に考えられる）

（取ったバッグに入っている金額の和 - 支払った金額 $A$ - 残った区間のDP結果）をよく考えると、

取ったバッグに入っている金額の和＝ (区間 $[l,r)$ の $X$ の和) - (残した区間の $X$ の和)

なので、残す区間を $[p,q)$ として、

(区間 $[l,r)$ の $X$ の和) - (区間 $[p,q)$ の $X$ の和) - $A$ - (区間 $[p,q)$ のDP結果)

「区間 $[l,r)$ の $X$ の和」と「支払った金額」は一定なので、結局

(区間 $[p,q)$ の $X$ の和) + (区間 $[p,q)$ のDP結果) … ①

を最小化するような $[p,q)$ を選べばよいことになる。

区間長 $w$ から $B$ 個取ったときに残るのは $w-B$ 個なので、$DP[w-B,(p=0,1,...)]$ から①をそれぞれ求めておく。（$Y_p$ とする）

$[l,r)$ を求めるときに使うのは $p=l～l+B$ なので、$Y_l～Y_{l+B}$ の最小値が求められればよい。

この区間最小値は、$l+1,l+2,...$ のDPを求めるときにも使い回したいことを考えると、スライド最小値などで求められる。

スライド最小値だと、$w$ ごとに $O(N)$ なので、全体で $O(N^2)$。セグメント木に載せてもよくて、$O(N^2 \log{N})$ となる。

「1つずつずれた区間の値を示す数列の、さらに区間最小値」を求めるので、添字など混乱してしまう。頑張る。

Python3

import os
import sys

import numpy as np


def solve(inp):
    n, a, b, c, d = inp[:5]
    xxx = inp[5:]

    acc = np.zeros(n + 1, np.int64)
    acc[1:] = xxx
    for i in range(n):
        acc[i + 1] += acc[i]

    dp = np.zeros((n + 1, n), np.int64)  # dp[w, l]: 長さ w 左端 l の区間からはじめて、先手が獲得できるスコア
    dp[1, :] = xxx

    for w in range(2, n + 1):

        # slide-min
        sm_ab_val = []
        sm_cd_val = []
        sm_ab = []
        sm_cd = []
        sm_ab_i = 0
        sm_cd_i = 0

        if w > b:
            x = w - b
            for i in range(n - x + 1):
                val = dp[x, i] + acc[i + x] - acc[i]
                sm_ab_val.append(val)
            for i in range(b + 1):
                while sm_ab and sm_ab_val[sm_ab[-1]] >= sm_ab_val[i]:
                    sm_ab.pop()
                sm_ab.append(i)

        if w > d:
            x = w - d
            for i in range(n - x + 1):
                val = dp[x, i] + acc[i + x] - acc[i]
                sm_cd_val.append(val)
            for i in range(d + 1):
                while sm_cd and sm_cd_val[sm_cd[-1]] >= sm_cd_val[i]:
                    sm_cd.pop()
                sm_cd.append(i)

        sm_ab_val.append(0)  # 最後の（無駄な）更新のためのダミー
        sm_cd_val.append(0)

        for l in range(n - w + 1):
            r = l + w

            res = max(-dp[w - 1, l + 1] + xxx[l], -dp[w - 1, l] + xxx[r - 1])

            if w <= b:
                res = max(res, acc[r] - acc[l] - a)
            else:
                tmp = acc[r] - acc[l] - sm_ab_val[sm_ab[sm_ab_i]] - a
                res = max(res, tmp)

                x = w - b
                if sm_ab[sm_ab_i] == l:
                    sm_ab_i += 1
                new_val = sm_ab_val[r - x + 1]
                while len(sm_ab) > sm_ab_i and sm_ab_val[sm_ab[-1]] >= new_val:
                    sm_ab.pop()
                sm_ab.append(r - x + 1)

            if w <= d:
                res = max(res, acc[r] - acc[l] - c)
            else:
                tmp = acc[r] - acc[l] - sm_cd_val[sm_cd[sm_cd_i]] - c
                res = max(res, tmp)

                x = w - d
                if sm_cd[sm_cd_i] == l:
                    sm_cd_i += 1
                new_val = sm_cd_val[r - x + 1]
                while len(sm_cd) > sm_cd_i and sm_cd_val[sm_cd[-1]] >= new_val:
                    sm_cd.pop()
                sm_cd.append(r - x + 1)

            dp[w, l] = res

    return dp[n, 0]


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)
    np.set_printoptions(linewidth=1000)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

Ex - Constrained Tree Degree

問題

$N$ 頂点の木のうち、以下の条件を満たすものの個数を $\mod{998244353}$ で答えよ
- 全ての頂点は区別できる（ラベル付き木）
- どの頂点の次数も、与えられる $K$ 個の整数集合 $S=\{S_1,...,S_K\}$ に含まれる
$2 \le N \le 2 \times 10^5$

解法

知識寄りの問題。

ラベル付き木は、プリューファーコードと呼ばれる数列に変換できることが知られている。

プリューファー列

数列の個数を数え上げることで、木の個数も数え上げられる。

頂点 $i$ の次数を $d_i$ とすると、$i$ はプリューファーコードにおいて必ず $d_i-1$ 回現れる。
（葉から処理する過程で、隣接する葉が消される時に自身が記録され、自身が葉になり処理される最後の1回は記録されないため）

プリューファーコードの長さは $N-2$ なので、

各 $i$ につき、$d_i \in S$
$(d_1-1)+(d_2-1)+...+(d_N-1)=N-2$

の2つを満たすような全ての $(d_1,...,d_N)$ の組み合わせについて、

1 を $d_1-1$ 個、2を $d_2-1$ 個、、、$N$ を $d_N-1$ 個並べるときの並べ方の個数は？
→ $\displaystyle \frac{(N-2)!}{(d_1-1)!(d_2-1)!...(d_N-1)!}$

を足し合わせたものが答えとなる。

分母のような形の計算は形式的冪級数を使うと表現しやすく、たとえば $S=\{1,4,6\}$ のとき、

$X = \dfrac{1}{0!}x^0 + \dfrac{1}{3!}x^3 + \dfrac{1}{5!}x^5$

という1頂点当たりの情報をもった $X$ を想定すると、$X^N$ における $x^{N-2}$ の係数に

$\displaystyle \sum_{条件を満たす全ての (d_1,...,d_N)}\frac{1}{(d_1-1)!(d_2-1)!...(d_N-1)!}$

が現れるので、$(N-2)!$ をかけたものが答えとなる。

繰り返し二乗法の要領で畳み込みをおこなうことで、$O(N(\log{N})^2)$ で答えが求められる。

Python3

import numpy as np


def precompute_factorials(n, MOD):
    f = 1
    factorials = [1]
    for m in range(1, n + 1):
        f = f * m % MOD
        factorials.append(f)
    f = pow(f, MOD - 2, MOD)
    finvs = [1] * (n + 1)
    finvs[n] = f
    for m in range(n, 1, -1):
        f = f * m % MOD
        finvs[m - 1] = f
    return factorials, finvs


def convolve(f, g):
    fft_len = 1
    true_len = len(f) + len(g) - 1
    while fft_len < true_len:
        fft_len <<= 1

    Ff = np.fft.rfft(f, fft_len)
    Fg = np.fft.rfft(g, fft_len)
    Fh = Ff * Fg

    h = np.fft.irfft(Fh, fft_len)
    h = np.rint(h).astype(np.int64)

    return h[:true_len]


def convolve_mod(f, g, p):
    f1, f2 = np.divmod(f, 1 << 15)
    g1, g2 = np.divmod(g, 1 << 15)

    a = convolve(f1, g1) % p
    c = convolve(f2, g2) % p
    b = (convolve(f1 + f2, g1 + g2) - (a + c)) % p

    h = (a << 30) + (b << 15) + c
    return h % p


n, k = map(int, input().split())
sss = list(map(int, input().split()))
MOD = 998244353

facts, finvs = precompute_factorials(n, MOD)

m = max(sss)
x = np.zeros(m + 1, np.int64)
for s in sss:
    x[s - 1] = finvs[s - 1]

loop = n
res = np.ones(1, np.int64)
while loop:
    if loop & 1:
        res = convolve_mod(res, x, MOD)[:n]
    x = convolve_mod(x, x, MOD)[:n]
    loop >>= 1

ans = res[n - 2] * facts[n - 2] % MOD
print(ans)

目次

NS Solutions Corporation Programming Contest 2023（AtCoder Beginner Contest 303）F,G,Ex問題メモ

F - Damage over Time

問題

解法

判定問題

解法2

G - Bags Game

問題

解法

高速化

Ex - Constrained Tree Degree

問題

解法