AtCoder Beginner Contest 217 G,H問題メモ

AtCoder Beginner Contest 217

G - Groups

• $1～N$ の番号が付いた人を空でない $k$ 個のグループに分ける
• このとき、番号を $M$ で割ったあまりが等しい人は同じグループになってはいけない
• $k=1～N$ のそれぞれについて、分け方の個数を $\mod{998244353}$ で求めよ
• $2 \le M \le N \le 5000$

同じグループになってはいけない人々の集合を、「不仲組」と呼ぶことにする。

不仲組は、$k$ がいくつの時でも $N,M$ だけで決まっている。

$N$ を $M$ で割った商とあまりを $d,m$ として、

• ① $d+1$ 人からなる不仲組が $m$ 組
• ② $d$ 人からなる不仲組が $M-m$ 組

N = 17   M = 4
不仲組1  1  5  9 13 17    d=4  m=1
不仲組2  2  6 10 14       5人の不仲組が1組
不仲組3  3  7 11 15       4人の不仲組が3組
不仲組4  4  8 12 16

不仲組毎に分け方を決めていく。
各不仲組のグループの分け方は他の不仲組に影響しないので、独立に決められる。

「ちょうど $k$ 個」は考えづらいので、「$k$ 個以下」のグループに分ける方法の個数を考える。

あらかじめ区別できる箱が $k$ 個並んでいて、そこに番号のついたボールを入れていくとすると、

• ① $k$ 個の箱に $d+1$ 個のボールを被らないように入れることを $m$ 回繰り返す… $({}_kP_{d+1})^m$
• ② $k$ 個の箱に $d$ 個のボールを被らないように入れることを $M-m$ 回繰り返す… $({}_kP_{d})^{M-m}$

この積が、全体を $k$ 個以下の箱に分ける方法の個数である。

ここから重複を除き、ちょうど $k$ 個の場合を求める。以下の2つが余計に数えられている。

• $k$ 個をフルに使っていないもの
• 箱を並べ替えたら同じになるもの

$k$ 個のうち、$l$ 個（$1 \le l \lt k$）しか使っていないものを考える。「ちょうど $l$ 個」の答え $Ans_l$ は既に求まっているとする。

$k$ 個の箱のどれをどの順に $l$ 個に対応づけるかで ${}_kP_l$ 通りなので、$l$ 個しか使っていないものは ${}_kP_l \times Ans_l$ 通りとなる。
$l=1～k-1$ まで走査して除外する。

残りは、$k$ 個をフルに使っているが箱を並べ替えたら同じになるものなので、$k!$ で割ればよい。

計算量は $O(N^2)$ となる。

Python3

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


n, m = map(int, input().split())
MOD = 998244353
facts, finvs = precompute_factorials(n + 5, MOD)

dd, md = divmod(n, m)
st = dd + (md > 0)

ans = [0] * (n + 1)
for k in range(st, n + 1):
    res = 1
    if md:
        b1 = facts[k] * finvs[k - dd - 1] % MOD
        b1 = pow(b1, md, MOD)
        res *= b1
    b2 = facts[k] * finvs[k - dd] % MOD
    b2 = pow(b2, m - md, MOD)
    res = res * b2 % MOD
    ans[k] = res

for i in range(st, n + 1):
    for j in range(st, i):
        ans[i] -= facts[i] * finvs[i - j] % MOD * ans[j] % MOD
    ans[i] *= finvs[i]
    ans[i] %= MOD

print('\n'.join(map(str, ans[1:])))

重複を除く部分は包除原理を適用することでも可能。

先ほどの「区別できる $k$ 個の箱に、空の箱を許して配る方法の個数」を $g(k)$ とする。
また、「区別できる“ちょうど” $k$ 個の箱に配る方法の個数」を $h(k)$ とする。
すると、求めたい答え $f(k)=\frac{h(k)}{k!}$ となる。

「$k$ 個の箱のうち、少なくとも $i$ 個が空である状態」という風に考えて包除原理を適用すると、 「$g(k-i)$ に対して新たに空の箱を $i$ 個追加して $k$ 個にするパターン」を引いたり足したりすればいいので、

• $h(k)=g(k) - g(k-1) {}_kC_{1} + g(k-2) {}_kC_{2} - ... g(0) {}_kC_{k}$
• $\displaystyle h(k)=\sum_{i=0}^{k} (-1)^i g(k-i) {}_kC_{i}$

これを変形すると、

• $\displaystyle h(k)=\sum_{i+j=k} (-1)^i g(j) {}_kC_{i}$
• $\displaystyle h(k)=\sum_{i+j=k} (-1)^i g(j) \frac{k!}{i!j!}$
• $\displaystyle h(k)=k! \sum_{i+j=k} \frac{(-1)^i}{i!} \frac{g(j)}{j!}$
• $\displaystyle f(k)=\sum_{i+j=k} \frac{(-1)^i}{i!} \frac{g(j)}{j!}$

となり、畳み込みに適した形となる。これで $O(N \log{N})$ で計算できる。

H - Snuketoon

• すぬけ君は時刻 $0$ に、無限に続く数直線上の座標 $0$ にいる
• すぬけ君は速さ1以下で数直線上を動ける
• 今から $N$ 回、左右の無限遠から水鉄砲が発射され、位置によってはダメージを負う
• $i$ 回目の発射
  • 時刻 $T_i$
  • 左右どちらから発車されるか $D_i$
  • どこまで届くか $X_i$
    • 左からの場合、すぬけ君が $X_i$ より左の座標 $p$ にいたら $|X_i-p|$ のダメージ
    • 右からの場合、すぬけ君が $X_i$ より右の座標 $p$ にいたら $|X_i-p|$ のダメージ
• すぬけ君が負うダメージの最小値を求めよ
• $1 \le N \le 2 \times 10^5$
• $1 \le T_i \le 10^9$ で、互いに相異なる

slope trickという手法を用いる。
名前は聞き馴染みが無かったが、図で考えると発想は理解しやすい。

公式editorialにもある通り、日本語資料としては以下の記事が図入りでわかりやすい。

• slope trick (1) 解説編 | maspyのHP

愚直には以下のDPを考えたくなる。（実際にはこの形では実装しない。$x$ の範囲がでかすぎるし）

• $DP[t][x]=$ 時刻 $t$ に座標 $x$ にいるときの最小コスト

どの $t$ の時点でも、$f(x)=DP[t][x]$ は「傾きが単調増加で、必ず整数」の形となる。
このように区間毎に直線に分割される関数を「区分線形関数」と称する。

傾き -3   -2  -1  0   1   4
       \                    /
        \                  /
         ＼               /
           ＼            /
             ~~--,__.--~~
    ----|---|----|--|----|-----
        2   5    8  10  12  

「傾きが単調増加で全て整数な区分線形関数 $f(x)$」は、以下の3つの情報を持っておけば復元できる。

• $f(x)$ の最小値
• 正の傾きの屈曲点の優先度付きキュー [ 10 12 12 12]
• 負の傾きの屈曲点の優先度付きキュー [ 2 5 8]

傾きが1増加する点を「屈曲点」と呼ぶことにする。
傾き $1→4$ のように一足飛びで変化する点は、複数の屈曲点が重なっていると見做す。

この3つを管理することで、$f(x)$ を更新していく。

主に以下の2つの方法となる。

• 次に水鉄砲が撃たれるまで移動する分の更新
• 水鉄砲が撃たれた分の更新

移動は、次に水鉄砲が撃たれるまでの時刻を $\Delta t$ として、最小値を中心に正負それぞれ $\Delta t$ だけ広げた形となる。

Δt = 3
       \                          /
        \                        /
         ＼                     /
           ＼                  /
             ~~--,________.--~~
    ----|---|----|--------|----|-----
       -1   2    5       13   15

• [ 10 12 12 12] → [ 13 15 15 15]
• [ 2 5 8] → [ -1 2 5]

キュー内の値全体に $\Delta t, -\Delta t$ が足される。
本当に全要素に加算すると時間かかるので、実際には外部でオフセットとして持っておく。

水鉄砲で受けるダメージは、以下のいずれかのように単純な関数の足し合わせとなる。

• $f(x)$ に $g(x)=\max(0, a-x)$ を足す
  • ある $x=a$ までは傾き $-1$ の直線で、そこから $0$ という形 ＼＿＿ をしている
• $f(x)$ に $g(x)=\max(0, x-a)$ を足す
  • ある $x=a$ までは $0$ で、そこから傾き $1$ の直線という形 ＿＿／ をしている

$a=X_i$ と、傾きが0の区間（負の屈曲点キューの最大値～正の屈曲点キューの最小値）の位置関係によって場合分けする。 詳細は冒頭の解説記事参照。

主に以下の操作で更新が可能となる。

• 正または負の屈曲点キュー（位置関係によって決まる）に、$a$ を追加する
• 最小値の区間が変化する場合、
  • 正から負の屈曲点キューへ（またはその逆へ）、先頭の要素を1つだけ移す
  • 最小値を更新する

傾き1や-1の関数を足すという操作では最小値をとる区間は高々1つしかズレないため、 屈曲点キュー間のやりとりも高々1回で済む、というのがポイント。

更新が1回 $O(\log{N})$ でできるため、全体で $O(N\log{N})$ となる。

この問題では初期状態は必ず座標0からスタートする制約がある。

一般的なslope trickでは $f(x)=0$ からスタートするため、 たとえば $t=1$ に座標 $10^9$ 未満に水鉄砲が撃たれても「はじめ $10^9$ にいたことにすればいいじゃん」みたいになりコストが狂ってしまう。

これは、以下の2つの方法で回避できる。

初期状態を、傾き0の区間が $[0,0]$ であり、両端の傾きが∞であると見做せばよい。

「傾き∞」を表現するには屈曲点キューに大量の $0$ を入れておけばよい。
今回の場合、最低限、更新回数である $N$ 個の $0$ を両方のキューに入れておけばよい。

または、答えが狂うのは正のキューから（真値からオフセットが引かれた後での）正の値が取り出されたり、 負のキューから負の値が取り出されたときなので、更新時にそれが発生した場合は取り出さずに0として扱う、としてもよい。

最後にはどこに居てもいいので、クエリを逆順に処理する。

最終的に $f(0)$ の値が答えとなる。
$f(0)$ を求めるのは $O(N)$ でできる。（キューから値を1個ずつ取り出して、傾き×区間長を足していけばいい）

Python3

from heapq import heappush, heapreplace

n = int(input())
neg_bend = [0] * n
pos_bend = [0] * n
neg_offset = 0
pos_offset = 0
current_min = 0
prev_t = 0

for _ in range(n):
    t, d, x = map(int, input().split())

    dt = t - prev_t
    neg_offset -= dt
    pos_offset += dt

    neg_head = -neg_bend[0] + neg_offset
    pos_head = pos_bend[0] + pos_offset

    if d == 0:
        if x <= pos_head:
            heappush(neg_bend, -(x - neg_offset))
        else:
            current_min += x - pos_head
            heappush(neg_bend, -(pos_head - neg_offset))
            heapreplace(pos_bend, x - pos_offset)
    else:
        if x >= neg_head:
            heappush(pos_bend, x - pos_offset)
        else:
            current_min += neg_head - x
            heappush(pos_bend, neg_head - pos_offset)
            heapreplace(neg_bend, -(x - neg_offset))

    prev_t = t

print(current_min)

今回、傾きの変化が1ずつだったので、最小値の区間がずれても高々1個のため、2本の優先度付きキューで管理できた。

しかし、「同様に ＿＿／ の形を足すが、その傾きが任意の整数」という場合は、一度に最小値の区間が大きくズレることで、 優先度付きキューだと一方から一方へ移す操作が多くかかってしまう。

その場合、平衡二分探索木で管理できるらしい。

• Generalized-Slope-Trick | Luzhiled’s Library