AtCoder Regular Contest 160 A,C,D問題メモ

AtCoder Regular Contest 160

A - Reverse and Count

問題

$(1, 2, …, N)$ の順列 $A$ が与えられる
$A$ の区間 $[L, R]$ の要素を反転させた順列を $f(L, R)$ とする
$(L, R)$ を $1 ≤ L ≤ R ≤ N$ から選ぶ方法は $N(N + 1)/2$ 通りある
すべての $(L, R)$ に対する順列 $f(L, R)$ を列挙し、辞書順に並べ替えたとき、$K$番目の順列を求めよ
$1 \le N \le 7000$

解法

indexを $0 ≤ L ≤ R ≤ N-1$ になおして考える。

$L$ を $L=0$ から順番に試す。

辞書順なので兎にも角にも先頭に“1”を持ってくるのが一番小さくて、その方法は（最初から $A_0$ が“1”だった場合を除き）$L=0$、$R=(A_i=1 であるような i)$ とする1通りである。

 4 3 1 6 2 5
 L   R
[1 3 4]6 2 5

また、“2”にするのも（最初から $A_0$ が“2”だった場合を除き）1通りで、これが2番目に小さい順列となる。

 4 3 1 6 2 5
 L       R
[2 6 1 3 4]5

同様に考えると、先頭が $A_0$ 以外になる順列はそれぞれ1通りしか作れず、順位が確定する。

4 3 1 6 2 5    → f(L,R)は21通り

辞書順    
     1    (先頭が1) 確定
     2    (先頭が2) 確定
     3    (先頭が3) 確定
     4    (先頭が4)
     5    (先頭が4)
     :
    19    (先頭が4)
    20    (先頭が5) 確定
    21    (先頭が6) 確定

先頭が“4”になるのは、$L \ge 1$ か、または $L=R=0$ であり実質 $A$ から変わらない順列のいずれかとなる。

ともかく、順位 1～3, 20～21 が確定し、4～19 が未確定となった。

これを $L=1,2,3,...$ で再帰的におこなえば、徐々に未確定順位が絞られていく。

$L=1$ の時、$A_1=3$
- $i=1$ より右にある $A_i$ で、3より小さいのは2個、大きいのは2個
- → 未確定順位 4～19 のうち、先頭2個 4～5 と、末尾2個 18～19 が確定
- → 残る未確定順位は 6～17 となり、$L=2$ に進む

処理の途中で $K$ 番目が確定すればよし、確定しなければ次に進むことを繰り返す。

$L=N-1$ まで処理すると、最終的に未確定順位が $N$ 個分残る。
これは $L=R$ とした時のものであり、$K$ が未確定順位内にある場合は $A$ のままが答えとなる。

Python3

from operator import add


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 + 1, 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)


def solve(n, k, ppp):
    rem = FenwickTreeInjectable(n + 1, int, add)
    for i in range(1, n + 1):
        rem.add(i, 1)

    a = 1
    b = n * (n + 1) // 2
    for l in range(n):
        p = ppp[l]
        rem.add(p, -1)
        lower = rem.sum(p)
        upper = n - l - lower - 1

        if a + lower > k:
            rank = k - a
            qqq = ppp[l:]
            qqq.sort()
            target = qqq[rank]
            r = ppp.index(target)
            return ppp[:l] + list(reversed(ppp[l:r + 1])) + ppp[r + 1:]

        if b - upper < k:
            rank = (n - l - 1) - (b - k)
            qqq = ppp[l:]
            qqq.sort()
            target = qqq[rank]
            r = ppp.index(target)
            return ppp[:l] + list(reversed(ppp[l:r + 1])) + ppp[r + 1:]

        a += lower
        b -= upper

    return ppp


n, k = map(int, input().split())
ppp = list(map(int, input().split()))

ans = solve(n, k, ppp)
print(*ans)

C - Power Up

問題

正整数の $N$ 要素の多重集合 $A=\{A_1,...,A_N\}$ が与えられる
以下の操作を0回以上、好きなだけおこなえる
- $A$ に2個以上含まれる要素を選ぶ（$x$ とする）
- $A$ から $x$ を2個削除し、$x+1$ を1個追加する
操作後の $A$ としてあり得る多重集合の種類数を $998244353$ で割ったあまりで求めよ
$1 \le N \le 2 \times 10^5$
$1 \le A_i \le 2 \times 10^5$

解法

DP。割と素直な考え方で解けるが、計算量をちゃんと評価する部分に若干の難しさがある。
（とはいえ、「何となく大丈夫そう」で通ってしまうことも多いけど）

$x$ 2個が $x+1$ 1個になることを、「$x+1$ に1個繰り上がる」と表現することにする。

以下のDPで解ける。

$DP[i][j]=i$ に $j$ 個繰り上がる、$i-1$ 以下の要素だけに操作を行った場合の多重集合の種類数

例えば初期状態で“1”が7個あったら、$0～3$ 個まで“2”に繰り上げることができるので、

         0  1  2  3
DP[2] = [1, 1, 1, 1]

となる。次、加えて“2”が例えば5個あったら、$DP[2]$ の $j=0～3$ につき、$DP[3]$ の $0～\lfloor \frac{j+5}{2} \rfloor$ に遷移できる。

Aの初期状態に"2"が5個:

DP[2][0]  →  DP[3][0～2] のそれぞれに加算
DP[2][1]  →  DP[3][0～3] のそれぞれに加算
DP[2][2]  →  DP[3][0～3] のそれぞれに加算
DP[2][3]  →  DP[3][0～4] のそれぞれに加算

結果     0  1  2  3  4
DP[3] = [4, 4, 4, 3, 1]

加算部分は累積和で高速化できる。

これを、$i$ が $\max(A)$ を超えた上で、$|DP[i]|=1$ になるまで繰り返せば、$DP[i][0]$ が答え。

計算量見積もり

$i$ の範囲は $\min(A)～(\max(A)+\log{N})$
- $\log{N}$ は、$\max(A)$ 以降もまだしばらく繰り上がれる場合があるから
$j$ は繰り上がりの個数なので最大 $\frac{N}{2}$

なので単純に掛け合わせると $O(N(A_{\max}+\log{N}))$ となり、TLEとなる。
しかし、それはあり得ない $(i,j)$ ペアまで含めてしまっているからで、有効なものだけを考えると、状態はぐっと減る。

ざっくり見積もりとして、以下が利用できる。

$A$ の初期状態に $i$ がない場合が続くと、$DP[i]$ の長さは $i$ が1増えるたびに半分ずつになっていく。

初期状態に $i$ があるとそこで長さが追加されるが、その場合も以下のように考えると、

細かな誤差とかは無視して、
「繰り上がってきた分」と「Aの初期状態にあった追加分」の
DP上での行く末を分けて考える
（本来は±1の端数のため、明確に分けられるものでもないが、ざっくり）

         0 1 2 3 4 5
DP[2]  [ o o o o o o ]
         ↓   ,-----'
DP[3]  [ o o o|o o o o o o ]
         ↓ ,-'   ,-------'
DP[4]  [ o o|o o o|o o o ]
         ~~~                2から追加された
             ~~~~~          3から追加された
                   ~~~~~    4から追加された

結局、それぞれの区間の長さは約半分ずつになっていくのは変わらない。
ある長さ $n$ の区間が半分ずつになっていって、なくなるまでの長さの合計は約 $2n$ である。

DPに追加される長さの総計が $O(N)$ なので、それが一括であろうと分割されていようと、全ての区間が無くなるまでの $DP[i]$ の長さの合計は $O(N)$ で収まることになる。

ただ、$N=2,A=\{1,200000\}$ のようなケースなど、繰り上がりが発生しなくても上限までの $DP[i][0]$ は遷移させ続けないといけないので、計算量は $O(N + A_{\max})$ となる。

Python3

D - Mahjong

問題

正整数 $N,M,K$ が与えられる
以下の条件を満たす数列 $A=(A_1,...,A_N)$ の個数を $998244353$ で割ったあまりで求めよ
- 長さは $N$、総和は $M$
- 以下の操作のどちらかを選んで行うことを繰り返すことで、全ての項を $0$ にできる
  - 任意の項 $A_i$ から $K$ 引く
  - 任意の連続する $K$ 項 $A_i,...,A_i+K-1$ から $1$ ずつ引く
$1 \le K \le N \le 2000$
$1 \le M \le 10^{18}$

解法

制約付き重複組み合わせ。計算量を減らす発想も必要。

大まかな方針イメージ

まず、どちらの操作にしても総和が $K$ ずつ少なくなるので、$M \mod{K} \equiv 0$ 以外の場合は不可能。
可能な場合、$d=\frac{M}{K}$ として、$d$ 回、いずれかの操作をすることになる。

操作を「縦」「横」で表現することにする。

縦操作のパターンは、$i=1～N$ のそれぞれに対して、そこに何回操作をおこなえるかで決まる。
横操作のパターンは、$i=1～N-K+1$ のそれぞれに対して、そこを左端とする横操作を何回おこなえるかで決まる。

横操作の回数を $y$、縦操作を $d-y$ として、それぞれの重複組み合わせ ${}_NH_{d-y} \times {}_{N-K+1}H_{y}$ とかで計算できそう。。。

重複除外

しかし、これでは以下の2パターンが同じになり、重複して数えてしまう。

ある連続する $K$ 項に縦操作をそれぞれ1回ずつおこなえるもの
同じ $i$ に横操作を $K$ 回おこなえるもの

重複を防ぐため、「横操作は、1箇所に対して $K-1$ 回までしかおこなえない」とする。（「縦操作は連続する $K$ 箇所に行えない」としてもよいが、まぁ、横操作の方が考えやすそう）

以下、1箇所の上限が $K-1$ までと決まっている重複組み合わせを、${}_nH'_{r}$ で表す。

計算量を削減する

横操作の回数を固定し、その場合の横操作と縦操作のパターン数の積を計算したい。
横操作の回数の上限は、$N-K+1$ 箇所に上限 $K-1$ を配置するため、$(N-K+1)(K-1)$ となる。

横操作の回数を $y$ と固定して、
$\displaystyle \sum_{y=0}^{(N-K+1)(K-1)} {}_{N-K+1}H'_{y} \times {}_{N}H_{d-y}$

数え上げる式は立てられたが、これでは計算量が多すぎる。個数上限のある重複組み合わせの求め方と、$y$ に対するループの部分がネック。

個数上限付き重複組み合わせの求め方

個数上限付き重複組み合わせは、例えば以下の2通りの求め方がある。

DP
包除原理

DPなら、${}_nH'_r$ は以下のサイトのように $O((nの範囲) \times (rの範囲))$ で出せるが、

個数制限がある場合の重複組合せの総数を動的計画法で求める - sambaiz-net

今回の場合、$n$ の上限は $N+K-1$、$r$ は横操作の回数上限 $(N-K+1)(K-1)$ なので、$O(N^3)$ となりDPでは無理そう。

包除原理は、個数上限が一定であるときに特に使いやすくて、 $n$ 個のうちいくつが（明示的に）制約違反しているかを求めることで、1回当たり $O(N)$ で求められる。

$\displaystyle {}_{N-K+1}H'_{y} = \sum_{t=0}^{N-K+1} (-1)^t \times {}_{N-K+1}C_{t} \times {}_{N-K+1}H_{y-tK}$
- あらかじめ制約違反すると決める $t$ 個の選び方 × それ以外の配り方

横と縦をまとめる工夫

包除原理で、上限付き重複組み合わせを $O(N)$ で計算できるとはいっても、

横方向の回数ごとに場合分けすると、$O(N^2)$ 回繰り返す必要がある
${}_nC_{r}$ を求める際、$n,r$ が非常に大きいことがあり、事前計算ができないので、1回ずつ計算する必要がある
- 今回の制約では $n-r$ が $O(N)$ に収まるので、$\frac{n(n-1)...(r+1)}{(n-r)!}$ を $O(N)$ で計算する

ために、全体でまだ $O(N^4)$ かかってしまう。

実は、縦操作と横操作の回数は、場合分けする必要が無い。
回数の和が $d$ と決まっている2つの重複組み合わせなので、「縦操作 $N$ と横操作 $N-K+1$、計 $2N-K+1$ から $d$ 個を重複を許して選ぶ」とすると、まとめて一括で考えられる。

$\displaystyle \sum_{y} {}_{N-K+1}H_{y} \times {}_{N}H_{d-y} = {}_{2N-K+1}H_{d}$

この場合でも、「その中で横操作に対応する $N-K+1$ 個だけは、$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


def ncr(n, r):
    if n < len(facts):
        return facts[n] * finvs[r] * finvs[n - r] % MOD

    res = finvs[n - r]
    for a in range(n, r, -1):
        res *= a
        res %= MOD

    return res


def nhr(n, r):
    return ncr(n + r - 1, r)


def solve(n, m, k):
    d, rem = divmod(m, k)

    if rem != 0:
        return 0

    ans = 0

    for t in range(min(d // k, n - k + 1) + 1):
        tmp = nhr(n + n - k + 1, d - t * k)
        tmp *= ncr(n - k + 1, t)
        if t & 1 == 0:
            ans += tmp
        else:
            ans -= tmp
        ans %= MOD

    return ans


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

facts, finvs = precompute_factorials(n * 2 + 5, MOD)

ans = solve(n, m, k)
print(ans)

目次

AtCoder Regular Contest 160 A,C,D問題メモ

A - Reverse and Count

問題

解法

C - Power Up

問題

解法

計算量見積もり

D - Mahjong

問題

解法

大まかな方針イメージ

重複除外

計算量を削減する

個数上限付き重複組み合わせの求め方

横と縦をまとめる工夫