AtCoder Beginner Contest 247 E,G,Ex問題メモ

AtCoder Beginner Contest 247

E - Max Min

問題

長さ $N$ の数列 $A=(A_1,A_2,...,A_N)$ と、整数 $X,Y$ が与えられる
連続した区間 $[l,r]=(A_l,A_{l+1},...,A_r)$ を考える
このうち、「最大値が $X$ 、最小値が $Y$ 」である区間（ $l,r$ の決め方）の個数を求めよ
$2 \le N \le 2 \times 10^5$

解法

公式Editorialでは、尺取、包除原理が紹介されている。

indexを管理する方法。

まず、 $X=Y$ の場合は別に処理する。
区間は $X$ が連続している箇所からしか選べない。 $X$ が $n$ 個連続した箇所なら $\dfrac{n(n+1)}{2}$ 区間取れる。
ランレングス圧縮し、これを合計すればよい。

以下、 $X \gt Y$ とする。また、「 $X$ より大きかったり、 $Y$ より小さい値」を総称して $Z$ と表現する。

以下の3つを整理する。

x_idx: $X$ が出現するindexリスト
y_idx: $Y$ が出現するindexリスト
z_idx: $Z$ が出現するindexリスト

求めたいのは、 $Z$ を含まず、 $X$ と $Y$ を両方含む区間の個数。

ある位置 $i$ にある $X$ を含むと固定して、少なくとも $i$ から見て最も直近の前後いずれかの $Y$ は含む必要がある。
それさえ満たせば、 $Z$ が出現するか、数列の端に到達するまで、左右独立に区間を伸ばすことができる。

つまり、 $X$ ごとに以下を（存在すれば）求めて合計すればよい。

その $X$ に対して、直前、直後の $Y$ の位置を求め、
- $Z$ を含まず (直前の $Y$ , $X$ ) を含む区間数
- $Z$ を含まず ( $X$ , 直後の $Y$ ) を含む区間数

以降、説明の便宜上、数列の左右端 $A_0,A_{N+1}$ にも、 $Z$ が存在するとする。

たとえば、(直前の $Y$ , $X$ ) を含む区間なら、

$i$ より前を遡っていったとき、 $Z$ より先に $Y$ が出現する必要がある
それを満たせば、区間数は以下で計算できる
- 区間の左端候補: 直前の $Z$ ～直前の $Y$
- 区間の右端候補: $X$ ～直後の $Z$
- 区間数 = 左端候補数 × 右端候補数

          #   *                     * 注目中のX
... Z _ _ Y _ X _ _ Y _ _ X Z ...   # 直前のY
     |-|-|                         区間左端候補 3
               |-|-|-|-|-|-|       区間右端候補 7  → 3x7=21 区間とれる

( $X$ , 直後の $Y$ ) を含む区間も同様にして数えられる。

重複除外

複数の $X$ を含む区間は、それぞれの $X$ を固定したときに同じ区間を重複して数えてしまう。
これを避けるため、 $X$ を固定する際は「固定した $X$ が、区間内の $X$ の中で一番左である区間だけを数える」ことにする。

また、その中でも、左の $Y$ と右の $Y$ を両方含む区間は重複してしまうため、「( $X$ , 直後の $Y$ ) の区間を計算する場合は、直前の $Y$ を含まない区間だけを数える」ことにする。

これで重複をのぞける。

計算量

計算量は、 $i$ から直前・直後の $Y$ や $Z$ の位置を求めるのに、リスト y_idx、z_idxに対して二分探索を使うと $O(N \log{N})$ となる。

または最初のindexリストを作るときに、

y_idx_L[i]: $x_idx[i]$ のすぐ左にある $Y$ のindex
y_idx_R[i]: $x_idx[i]$ のすぐ右にある $Y$ のindex

みたいにしておくと、これは $O(N)$ で構築できるので、全体でも $O(N)$ となる。

Python3

G - Dream Team

問題

$N$ 人の競技プログラミングプレイヤーがいる
人 $i$ は、所属 $A_i$ 、得意分野 $B_i$ 、強さ $C_i$ が決まっている
「全ての構成員が、互いに異なる所属、異なる得意分野を持つチーム」を「ドリームチーム」と呼ぶ
- ドリームチームの強さは、構成員の強さの和
以下を求めよ
- $N$ 人の中から選んで作れるドリームチームの最大構成員数 $k$
- $i=1～k$ のそれぞれにつき、 $i$ 人のドリームチームの強さの最大値
$1 \le N \le 30000$
$1 \le A_i,B_i \le 150$

解法

最小費用流。ただし1回流して終わりではなく、追加で何回か流す。

グラフを作る

所属と得意分野の種類数が150と少ない。これを頂点としたフローでも計算量的に大丈夫そう。

      A     B
 ,--->①--->①---,
 |      `-_      |
Ⓢ--->②   '②---+-> Ⓣ
 |       ,-'     |
 |--->③'   ③---|
 :    ：    ：   :

Ⓢから $A$ の各頂点へ、容量1、コスト0の辺
各プレイヤーにつき、 $A_i$ から $B_i$ へ、容量1、コスト $X - C_i$ の辺
- 今回は最大化の問題なので、最小費用流の上では正負逆転させたいが、負辺を防ぐため、適当にでかい数字 $X$ を決め、コストを $X - C_i$ として解く
$B$ の各頂点からⓉへ、容量1、コスト0の辺

こうすれば、各 $A_i,B_i$ を経由する流量は上限1に限定され、複数の流れが同じ $A_i,B_i$ を通ることがなくなる。
つまり、流量 $f$ 流せば、 $f$ 本の流れはすべて別々の $A_i,B_i$ を通り、ドリームチームの条件を満たす。

またそのコストはチームメンバーの $X - C_i$ の合計値となる。
「メンバーの人数 × $X$ - 最小費用流の結果」とすれば、 $C_i$ の合計に戻せる。

流量 $f$ が流せなくなったら、その直前 $f-1$ が、 $N$ 人で作れるドリームチームの最大構成員数 $k$ となる。

複数回流す

流量 $3$ 流せば3人のドリームチームの強さの最大値が求まり、 $4$ 流せば4人の（略）が求まる。

が、毎回新規にグラフを構築したりリセットしたりして、流量を $1,2,3,...$ と増やしながら調べるのは非効率。

最小費用流は、たとえば流量 $f$ を流し終えた後のグラフに対し、「さらに追加で $g$ 流す」ことができる。
この場合、流量 $f$ の結果からの増加分が結果となる。

なので、グラフは再構築やリセットせず、1ずつ追加で流すことを繰り返せばよい。

Python3

Ex - Rearranging Problem

問題

$N$ 人の人 $1,2,...,N$ が順に並んでいて、人 $i$ は色 $c_i$ の服を着ている
「任意の2人を選んでswapする」という操作をちょうど $K$ 回繰り返した
$K$ 回の操作後も、操作前と服の色の並び順は同じだった
- つまり、全ての $i$ について、先頭から $i$ 番目の人は色 $c_i$ の服を着ていた
このとき、操作後の人の並び順として考えられるものが何通りか、 $\mod{998244353}$ で求めよ
$2 \le N \le 2 \times 10^5$

解法

求める並び順の性質

操作回数の偶奇

まず、順列（のように同じ要素を含まない数列）の転倒数を考えると、1回のswapによっては必ず奇数しか変化しない。

... 3 1 5 9 7 ...

3と7をswapする場合、その間にある数字は以下のいずれかだが、

3未満かつ7未満
3より7より大きい
3と7の間

どれをとっても、その数字と3,7との転倒数は「2減る」「2増える」「1増えて1減って、相殺されて0」のいずれかしかない。
また、3と7自体の転倒数は、swapによって1増えるか1減る。よって、全体では奇数の変化となる。

初期状態（転倒数0＝偶数）から、奇数回の操作後は転倒数は奇数、偶数回なら偶数となる。
奇数回操作しても偶数回操作しても作れる並び順、というものは存在しない。

また、 $K$ 回より少ない回数で達成できる並び順で $K$ と偶奇が同じものは、適当に同じ要素に2回ずつ操作すれば回数を稼げるので、ちょうど $K$ 回でも達成可能となる。

よって、「その並び順を作るための最小操作回数」ごとに並び順の個数を数えることができれば、そこから「 $K$ 以下」かつ「 $K$ と偶奇が一致する」ものの総和を取ることで、答えとなる。

swapは同色同士だけ考えればOK

「初期状態」と「 $K$ 回の操作後」で色の並び順が同じならばよくて、途中経過は違っていてもよい。
とはいうもののこれは若干ミスリードっぽくて、今は「最小操作回数」を考えたいので、違う色をswapするのは無駄となる。

同じ色の中だけでswapすると考えてよい。

最小操作回数別に数える

同色の中だけでswapが行われるとして、最小操作回数別に並び順の個数を数えたい。

よく、「順列 $P$ が与えられる。swapによって $1,2,...,N$ を $P_1,P_2,...,P_N$ に並べ替える最小操作回数を求めよ」みたいな問題がある。
これには、サイクルに分割する解法がよく用いられる。
サイクル毎に、サイクルサイズ-1の回数を要するので、その総和が答え。

ということは全体で考えると、長さ1のサイクルもサイクルとして考えて、 $N$ 個の数が「いくつのサイクルに分割されたか」によって $N-(サイクル数)$ が答えともいえる。

そして、 $N$ 個の数を、 $K$ 個のサイクルに（サイクル内の並び順も区別して）分ける方法の個数 $S_{n,k}$ は、第1種スターリング数として知られている。

スターリング数

例えばある色の服を着た人が $m$ 人いたら、最小操作回数毎の並び順の個数は、以下のようになる。

$m$ 個のサイクルに分ける方法（最小操作回数 $0$ ）: $S_{m,m}$
$m-1$ 個のサイクルに分ける方法（最小操作回数 $1$ ）: $S_{m,m-1}$
：
$1$ 個のサイクルに分ける方法（最小操作回数 $m-1$ ）: $S_{m,1}$

各色をまとめる

色	人数	最小操作回数
色	人数	0	1	…	$a-1$	…	$b-1$
色1の服を着た人	$a$	$S_{a,a}$	$S_{a,a-1}$		$S_{a,1}$
色2の服を着た人	$b$	$S_{b,b}$	$S_{b,b-1}$		$S_{b,b-a+1}$		$S_{b,1}$

というのが仮に計算できたとして、ではこの2色を同時に考慮した最小操作回数毎の並び順はどうなるかというのは、これらの畳み込みによって表現できる。

	最小操作回数
	$0$	$1$	…	$b-a+2$
色1の服を着た人	$S_{a,a}$	$S_{a,a-1}$
色2の服を着た人	$S_{b,b}$	$S_{b,b-1}$
2色を合わせて考慮	$S_{a,a}S_{b,b}$	$S_{a,a-1}S_{b,b}+S_{a,a}S_{b,b-1}$		$S_{a,1}S_{b,1}$

他の色も、どんどん畳み込んでまとめていけばよい。
最終的に1つになった数列が、全体での「最小操作回数別の並び順の個数」となる。

畳み込み後の長さ＝全体の最小操作回数の最大値は $N-1$ なので、長くなりすぎて計算量的にアウト、ということはない。

スターリング数を求める

さて、先ほど、同色の人が $m$ 人いる場合の最小操作回数毎の並び順 $[S_{m,m},S_{m,m-1},...,S_{m,1}]$ は計算できた前提だったが、ここも問題がある。

スターリング数はDP的に求めるしか無く（多分）、一般的な方法では $S_{m,○}$ を求めるには $O(m^2)$ かかる。
$m$ は最大 $N$ になり得るので、これだとTLE。

しかし、スターリング数にはもう一つ、便利な表現の仕方があって、

$x(x+1)(x+2)...(x+n-1)$ を展開したときの、 $x^k$ の係数が $S_{n,k}$ となる

従って、ここでも畳み込みが使える。

今回は、スターリング数の $n$ 段目を逆並びにしたものを求めたいので、同色の人数 $m$ に対して

$[1,1],[1,2],[1,3],...,[1,m-1]$

を全て畳み込むと、求めたい数列となる。

畳み込みで結合するに従って数列サイズはだんだん肥大化する。
小さい方から結合するようにしていく（マージテク）ことによって、計算量を抑えられる。

Python3

numpy.fftを使った畳み込みではTLEが取り切れなかった。

（numpy.fftは小数で演算されるため、オーバーフローする可能性のある畳み込みでは、bitをずらして3回行う必要がある）

まぁ、それ以外にもよろしくない実装があったのかも知れないが、畳み込みの処理を直書きするのはやはり重たい。

        
          
          
              
              import os
import sys
from heapq import heapify, heappop, heapreplace
 
import numpy as np
 
 
def solve(inp):
    def bit_length(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x
 
    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x
 
    def pow_mod(x, n, m):
        r = 1
        y = x % m
        while n:
            if n & 1:
                r = (r * y) % m
            y = y * y % m
            n >>= 1
        return r
 
    def get_primitive_root(m):
        if m == 2: return 1
        if m == 167772161: return 3
        if m == 469762049: return 3
        if m == 754974721: return 11
        if m == 998244353: return 3
        divs = [2]
        x = (m - 1) // 2
        while x & 1 == 0:
            x >>= 1
        i = 3
        while i * i <= x:
            if x % i == 0:
                divs.append(i)
                x //= i
                while x % i == 0:
                    x //= i
            i += 2
        if x > 1:
            divs.append(x)
        g = 2
        while True:
            ok = True
            for d in divs:
                if pow_mod(g, (m - 1) // d, m) == 1:
                    ok = False
                    break
            if ok:
                return g
 
    def butterfly_prepare(mod, primitive_root):
        sum_e = np.zeros(30, np.int64)
        sum_ie = np.zeros(30, np.int64)
        es = np.zeros(30, np.int64)
        ies = np.zeros(30, np.int64)
        cnt2 = bit_scan_forward(mod - 1)
        e = pow_mod(primitive_root, (mod - 1) >> cnt2, mod)
        ie = pow_mod(e, mod - 2, mod)
        for i in range(cnt2, 1, -1):
            es[i - 2] = e
            ies[i - 2] = ie
            e = e * e % mod
            ie = ie * ie % mod
        now_e = 1
        now_ie = 1
        for i in range(cnt2 - 1):
            sum_e[i] = es[i] * now_e % mod
            sum_ie[i] = ies[i] * now_ie % mod
            now_e = now_e * ies[i] % mod
            now_ie = now_ie * es[i] % mod
        return sum_e, sum_ie
 
    def butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = bit_length(n)
 
        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod
 
    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = bit_length(n)
 
        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod
 
    def convolve_ntt(aaa, bbb, MOD):
        n = aaa.size
        m = bbb.size
        k = n + m - 1
        z = 1 << bit_length(k)
        raaa = np.zeros(z, np.int64)
        rbbb = np.zeros(z, np.int64)
        raaa[:n] = aaa
        rbbb[:m] = bbb
        primitive_root = get_primitive_root(MOD)
        sum_e, sum_ie = butterfly_prepare(MOD, primitive_root)
        butterfly(raaa, MOD, sum_e)
        butterfly(rbbb, MOD, sum_e)
        for i in range(z):
            raaa[i] = raaa[i] * rbbb[i] % MOD
        butterfly_inv(raaa, MOD, sum_ie)
        iz = pow_mod(z, MOD - 2, MOD)
        for i in range(k):
            raaa[i] = raaa[i] * iz % MOD
        return raaa[:k]
 
    MOD = 998244353
 
    k = inp[1]
    ccc = inp[2:]
 
    counter = {}
    for c in ccc:
        if c in counter:
            counter[c] += 1
        else:
            counter[c] = 1
 
    targets = []
    hq = []
 
    for cnt in counter.values():
        for i in range(1, cnt):
            targets.append(np.array([1, i], dtype=np.int64))
            hq.append((2, len(targets) - 1))
 
    if len(hq) == 0:
        if k % 2 == 0:
            return 1
        else:
            return 0
 
    heapify(hq)
 
    while len(hq) > 1:
        c1, i1 = heappop(hq)
        c2, i2 = hq[0]
        targets[i2] = convolve_ntt(targets[i1], targets[i2], MOD)
        heapreplace(hq, (c1 + c2 - 1, i2))
 
    _, i = hq[0]
    pat = targets[i]
 
    if k % 2 == 0:
        ans = pat[0:k + 1:2].sum() % MOD
    else:
        ans = pat[1:k + 1:2].sum() % MOD
    return ans
 
 
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)

            

        
      

−目次

AtCoder Beginner Contest 247 E,G,Ex問題メモ

E - Max Min

問題

解法

重複除外

計算量

G - Dream Team

問題

解法

グラフを作る

複数回流す

Ex - Rearranging Problem

問題

解法

求める並び順の性質

操作回数の偶奇

swapは同色同士だけ考えればOK

最小操作回数別に数える

各色をまとめる

スターリング数を求める