第5回 ドワンゴからの挑戦状 本選 A - Taro vs. Jiro

このコンテストの作問チームにいたが、この問題は実装しなかったので今更実装した。(は?)
A - Taro vs. Jiro

問題

2人が、頂点に赤青の色が塗られた有向グラフの上にコマを置いてゲームをする。各プレイヤーは1ターンに1度、コマを置かれた頂点に隣接する頂点のどれかに動かす。Kターン後にコマが置かれた頂点の色が青なら先手の勝ち、赤なら後手の勝ち。各頂点について、そこからゲームを開始した時、どちらが勝つか判定せよ。

  • 2 <= (頂点数) <= 10^5
  • 1 <= (辺数) <= 10^5
  • 1 <= K <= 10^18
  • 与えられるグラフは単純で連結

解法

K>=3の場合、いずれのプレイヤーも、直前の相手の動きを打ち消すことができる。よって、K手での勝敗は(K-2)手での勝敗に一致する。
K=1のときとK=2のときに判定できればよい。
K=1のとき、初期地点をvとしたとき、vに隣接する頂点に青が一つでも存在する時、およびときに限り先手は勝てる。K=2のとき、初期地点をvとしたとき、vに隣接する頂点に(1手で赤に行ける頂点)が一つも存在しないとき、およびそのときに先手は勝てる。

実装上の注意点

  • K=2のとき、単に2手で行ける場所を全列挙してしまうとO(N^2)時間かかるので、DPなどをする必要がある。
  • 今回は意図的に除外されているが、K=0のケースがコーナーケースなので、入力としてあり得るかどうかは絶対確認する!

提出: Submission #3867920 - 第5回 ドワンゴからの挑戦状 本選 (Rust)

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
        m: usize,
        k: i64,
        s: chars,
        ab: [(usize1, usize1); m],
    }
    let mut g = vec![Vec::new(); n];
    for (a, b) in ab {
        g[a].push(b);
        g[b].push(a);
    }
    let mut taro = vec![false; n];
    if k % 2 == 1 {
        for v in 0 .. n {
            for &w in g[v].iter() {
                if s[w] == 'B' {
                    taro[v] = true;
                }
            }
        }
    } else {
        let mut dp = vec![[false; 2]; n];
        for v in 0 .. n {
            for &w in g[v].iter() {
                if s[w] == 'R' {
                    dp[v][0] = true;
                }
            }
        }
        for v in 0 .. n {
            for &w in g[v].iter() {
                if !dp[w][0] {
                    dp[v][1] = true;
                }
            }
        }
        for i in 0 .. n {
            taro[i] = dp[i][1];
        }
    }
    for ans in taro {
        puts!("{}\n", if ans { "First" } else { "Second" });
    }
}

まとめ

この問題は本戦参加者全員が解けるくらいの難易度を目指して作られた。ほとんどの人に解かれたようでよかった。(解いていない赤コーダーは初手CかDに突っ込んだはずなのでカウントしない。)

CADDi 2018 F - Square

線形代数をやりすぎて迷走してしまった。
https://atcoder.jp/contests/caddi2018/tasks/caddi2018_d

問題

N行N列のマス目がある。これに0または1の値を書き込む。マス目のうちいくつかにはすでに数字が書き込まれている。
完成したマス目は以下の条件を満たす必要がある。

  • 1 <= i < j <= Nであるような整数i,jに対して、(i,i)を左上、(j,j)を右下にするような正方形の中に含まれる1の個数は偶数である。

条件を満たすマス目の埋め方を数え上げよ。

解法

(i,j)に書き込む数字をa[i][j]と表記する。
まず、|i-j| >= 3ならばa[i][j] = a[j][i]である。これは冒頭の規則を(i,j), (i, j-1), (i+1,j), (i+1,j-1)に対して適用すればわかる。
残りは対角成分付近である。対角成分に着目すると、a[i+1][i+1] = a[i][i]+(a[i][i+1]+a[i+1][i]), a[i+1][i+1] = a[i][i+2] + a[i+2][i]がわかる。
ここから、対角成分付近の自由度を計算しつつ、ありえる対角成分の数え上げをすればよい。
対角成分の数え上げは、以下の問題を解くことで可能である。

  • 長さnの0または1からなる数列bおよびそのmod 2での階差数列cについて、b[i]=jという形の制約と、c[i]=jという形の制約が何個か与えられる。条件を全て満たす(b,c)を数え上げよ。

これはDPでできる。(コード中のcalc_acc)
また、対角成分付近の自由度の数え上げは、各iについて(i+1,i)と(i,i+1)についてはmax(0,もともと埋まっていない個数-1)が自由度、(i+2,i)と(i,i+2)についても同様、とすればよい。(マイナス1は上のcの数え上げ部分に吸収されている分であることに注意。)

実装上の注意点

  • 端で壊れがちなので気をつける
  • 制約の矛盾などを発見するときにビット演算を使うが、そこでバグらせないように細心の注意を払う
  • 数え上げだが、実質自由度を求める問題なので、本当に自由度があっているか確認する

提出: Submission #3849384 - CADDi 2018 (Rust)

// ModInt 省略
fn calc_acc(acc: &[i32], inb: &[i32]) -> ModInt {
    let n = acc.len();
    assert_eq!(inb.len(), n - 1);
    // eprintln!("acc = {:?}, inb = {:?}", acc, inb);
    let mut dp = vec![[ModInt::new(0); 2]; n];
    let one = ModInt::new(1);
    if acc[0] == 0 {
        dp[0][0] = one;
        dp[0][1] = one;
    } else {
        dp[0][acc[0] as usize - 1] = one;
    }
    for i in 1 .. n {
        if inb[i - 1] == 0 {
            for k in 0 .. 2 {
                dp[i][0] += dp[i - 1][k];
            }
            dp[i][1] = dp[i][0];
        } else {
            let diff = (inb[i - 1] - 1) as usize;
            for k in 0 .. 2 {
                dp[i][k ^ diff] += dp[i - 1][k];
            }
        }
        if acc[i] != 0 {
            let opp = 1 - (acc[i] - 1) as usize;
            dp[i][opp] = ModInt::new(0);
        }
    }
    dp[n - 1][0] + dp[n - 1][1]
}


fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
        m: usize,
        abc: [(i32, i32, i32); m],
    }
    let mut dim: i64 = 0;
    let mut hm = HashMap::new();
    for &(a, b, c) in abc.iter() {
        let (a, b) = if a < b || (a - b).abs() <= 2 { (a - 1, b - 1) } else { (b - 1, a - 1) };
        *hm.entry((a, b)).or_insert(0) |= 1 << c;
    }
    // rest, |x - y| <= 2, even
    let mut rest = vec![vec![0; 3]; n];
    let mut odd = vec![Vec::new(); n - 1];
    for (&(a, b), &val) in hm.iter() {
        if val == 3 {
            puts!("0\n");
            return;
        }
        if (a - b).abs() >= 3 {
            dim -= 1;
        }
        if (a - b).abs() <= 2 && (a + b) % 2 == 0 {
            let idx = (a + b) as usize / 2;
            let u = (b - a + 2) as usize / 2;
            rest[idx][u] = val;
        }
        if (a - b).abs() == 1 {
            odd[(a + b) as usize / 2].push((a, b, val));
        }
    }
    let mut adden = 0;
    if n > 3 {
        for i in 0 .. n {
            let lock = if i <= 1 {
                3 + i as i64
            } else if i <= n - 2 {
                5
            } else {
                (n - 1 - i) as i64 + 2
            };
            adden += n as i64 - lock;
        }
    }
    dim += adden / 2;
    // eprintln!("dim = {}", dim);
    let mut cons = vec![0; n]; // constraints
    for i in 0 .. n {
        cons[i] |= rest[i][1];
    }
    // rest even
    for i in 1 .. n - 1 {
        if rest[i][0] != 0 && rest[i][2] != 0 {
            let x = (rest[i][0] - 1) ^ (rest[i][2] - 1);
            cons[i] |= 1 << x;
        }
        let mut f = 0;
        if rest[i][0] == 0 { f += 1; }
        if rest[i][2] == 0 { f += 1; }
        dim += max(0, f - 1);
    }
    for i in 0 .. n {
        if cons[i] == 3 {
            puts!("0\n");
            return;
        }
    }
    let mut inb = vec![0; n - 1];
    for i in 0 .. n - 1 {
        if odd[i].len() == 2 {
            let x = (odd[i][0].2 - 1) ^ (odd[i][1].2 - 1);
            inb[i] |= 1 << x;
        }
        let mut f = 2 - odd[i].len() as i64;
        dim += max(f - 1, 0);
    }
    // rest odd
    let fac = calc_acc(&cons, &inb);
    /*
    eprintln!("dim = {}", dim);
    eprintln!("fac = {}", fac);
     */
    puts!("{}\n", ModInt::new(2).pow(dim) * fac);    
}

まとめ

  • 頭が数学モードになっていると、競技プログラミングのことをすっかり忘れてしまう…。
  • AtCoderのRustのバージョンが低いため最初の提出がCEになり、順位が1つ下がった (は?)

CADDi 2018 C - Product and GCD

これがやりたかっただけだろシリーズ。素因数分解のためにポラードのロー法を貼った。
https://atcoder.jp/contests/caddi2018/tasks/caddi2018_a

問題

N個の正の整数a_1, ..., a_Nについて、それらの積はPである。
このとき、gcd(a_1,...,a_N)としてあり得る値の最大値を求めよ。
1 <= N, P <= 10^12

解法

Pを素因数分解してp_i^e_iの積になったとする。p_i^floor(e_i/N)の積が答えである。
素因数分解は普通にやってもO(sqrt(P))くらいでできるが、ポラードのロー法などの速いアルゴリズムならO(P^{1/4})くらいでできるはず。

実装上の注意点

  • e_iはそこまで大きくならないので、p_i^floor(e_i/n)の積は単純に掛けていくだけでよい。

提出: Submission #3837935 - CADDi 2018 (Rust)

mod pollard_rho {
    use std::collections::HashMap;
    /// binary gcd
    pub fn gcd(mut x: i64, mut y: i64) -> i64 {
        if y == 0 { return x; }
        if x == 0 { return y; }
        let mut sh = 0;
        while ((x | y) & 1) == 0 {
            x >>= 1; y >>= 1; sh += 1;
        }
        while (x & 1) == 0 { x >>= 1; }
        while y != 0 {
            while (y & 1) == 0 { y >>= 1; }
            if x > y { let t = x; x = y; y = t; }
            y -= x;
        }
        x << sh
    }

    fn add_mod(x: i64, y: i64, n: i64) -> i64 {
        let z = x + y;
        if z >= n { z - n } else { z }
    }

    fn mul_mod(x: i64, mut y: i64, n: i64) -> i64 {
        assert!(x >= 0);
        assert!(x < n);
        let mut sum = 0;
        let mut cur = x;
        while y > 0 {
            if (y & 1) == 1 {
                sum = add_mod(sum, cur, n);
            }
            cur = add_mod(cur, cur, n);
            y >>= 1;
        }
        sum
    }

    fn mod_pow(x: i64, mut e: i64, n: i64) -> i64 {
        let mut prod = if n == 1 { 0 } else { 1 };
        let mut cur = x % n;
        while e > 0 {
            if (e & 1) == 1 {
                prod = mul_mod(prod, cur, n);
            }
            cur = mul_mod(cur, cur, n);
            e >>= 1;
        }
        prod
    }

    pub fn is_prime(n: i64) -> bool {
        if n <= 1 { return false; }
        let small = [2, 3, 5, 7, 11, 13];
        if small.iter().any(|&u| u == n) { return true; }
        if small.iter().any(|&u| n % u == 0) { return false; }
        let mut d = n - 1;
        let mut e = 0;
        while (d & 1) == 0 {
            d >>= 1;
            e += 1;
        }
        let a = [2, 325, 9375, 28178, 450775, 9780504, 1795265022];
        a.iter().all(|&a| {
            if a >= n { return true; }
            let mut x = mod_pow(a, d, n);
            if x == 1 { return true; }
            for _ in 0 .. e {
                if x == n - 1 {
                    return true;
                }
                x = mul_mod(x, x, n);
                if x == 1 { return false; }
            }
            x == 1
        })
    }

    fn pollard_rho(n: i64, c: &mut i64) -> i64 {
        if n % 2 == 0 { return 2; }
        loop {
            let mut x: i64 = 2;
            let mut y = 2;
            let mut d = 1;
            let cc = *c;
            let f = |i| add_mod(mul_mod(i, i, n), cc, n);
            while d == 1 {
                x = f(x);
                y = f(f(y));
                d = gcd((x - y).abs(), n);
            }
            if d == n {
                *c += 1;
                continue;
            }
            return d;
        }
    }

    /// Outputs (p, e) in p's ascending order.
    pub fn factorize(x: i64) -> Vec<(i64, usize)> {
        if x <= 1 {
            return Vec::new();
        }
        let mut hm = HashMap::new();
        let mut pool = vec![x];
        let mut c = 1;
        while let Some(u) = pool.pop() {
            if is_prime(u) {
                *hm.entry(u).or_insert(0) += 1;
                continue;
            }
            let p = pollard_rho(u, &mut c);
            pool.push(p);
            pool.push(u / p);
        }
        let mut v: Vec<_> = hm.into_iter().collect();
        v.sort();
        v
    }
} // mod pollard_rho
use pollard_rho::*;

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: i64,
        p: i64,
    }
    let fac = factorize(p);
    let mut ans = 1;
    for (a, e) in fac {
        let e = e as i64 / n;
        for _ in 0 .. e { ans *= a; }
    }
    puts!("{}\n", ans);
    
}

まとめ

ポラードのロー法で殴ったので記念に書いただけです。ライブラリとしてでも使ってください。

AGC029 D - Grid game

ゲーム(高橋君に選択権があるとは言ってない)みたいな感じ。
https://atcoder.jp/contests/agc029/tasks/agc029_d

問題

H行W列のマス目の上に障害物がN個とコマが1個(1,1)に置かれている。x行y列の位置を(x,y)と表記する。
高橋君と青木君がマス目の上でゲームをする。
高橋君が先手で、以下の行動を交互にする。

  • コマの位置が(x,y)のとき、高橋君は(x+1,y)に、青木君は(x,y+1)にコマを移動させる。あるいはコマを動かさない。

二人が連続してコマを動かさないことを選択したときにゲームは終わる。高橋君は自分が行う行動の回数を最大化することを、青木君は自分が行う行動の回数を最小化することを目標に動く。二人が最適に行動する時、高橋君の行動回数を求めよ。

  • 1 <= H, W <= 2 * 10^5
  • 0 <= N <= 2 * 10^5
  • 障害物は(1, 1)にはない

解法

高橋君は、自分が止まった瞬間青木君に止まられてゲームが終了してしまう。つまり、高橋君はできるならずっと動き続けないといけない。
青木君は高橋君を障害物に早くぶつけることを考えればいい。最小化すべきなのは高橋君の行動回数(<=> 高橋君がぶつかる障害物のx座標)なのだから、それを求めるためには、貪欲にy軸方向に進んでいって、通った場所からx軸の+側にある直近の障害物のx座標の最小値を求めればいい。
以上を高速にするためには、1からWまでのy座標ごとに、そこにある障害物のx座標を昇順にして配列などに持っておけばいい。

実装上の注意点

  • (x,y)という記号から暗黙に高橋君が右にコマを動かすと思いがちだが、そうではない

- それに付随して、縦横の間違いに気を付ける

  • ループを終了させるタイミングに気をつける(W回ではなくH回)

提出: Submission #3830128 - AtCoder Grand Contest 029 (Rust)

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        h: usize,
        w: usize,
        n: usize,
        xy: [(usize1, usize1); n],
    }
    let mut occur = vec![Vec::new(); w];
    for (x, y) in xy {
        occur[y].push(x);
    }
    for i in 0 .. w { occur[i].sort(); }
    let mut cx = 0;
    let mut cy = 0;
    let mem = |x, y: usize| occur[y].binary_search(&x).is_ok();
    let mut ans = h;
    for _ in 0 .. h {
        let idx = occur[cy].binary_search(&cx).unwrap_err();
        ans = min(ans, if idx >= occur[cy].len() { h } else { occur[cy][idx] });
        if cx >= h - 1 || mem(cx + 1, cy) {
            break;
        }
        cx += 1;
        if cy < w - 1 && !mem(cx, cy + 1) { cy += 1; }
    }
    puts!("{}\n", ans);
}

まとめ

C問題よりも実装が短いんですが…。

AGC029 C - Lexicographic constraints

本番は出場せず、Twitterで「二分探索」という単語を見てから解法を考えたため、自力で思いつけたかどうかは自信がない。
https://atcoder.jp/contests/agc029/tasks/agc029_c

問題

N個の文字列S_1, ..., S_nがある。S_iの長さはA_iであり、S_1 < S_2 < ... < S_nが成り立っている。このとき、あり得るアルファベットのサイズ(文字の種類数)として一番小さいものを求めよ。

  • 1 <= N <= 2 * 10^5
  • 1 <= A_i <= 10^9

解法

実現可能性はアルファベットのサイズについて単調である。つまり、アルファベットのサイズが増えることで可能だったものが不可能になることはない。
これを利用して二分探索を行うと、決められたアルファベットのサイズkについて、S_1, ..., S_nとしてあり得るものが存在するか判定すれば良いことがわかる。
先頭から貪欲に小さい文字列を割り当てていくのが最善なので、A_i < A_{i+1}ならS_iに'a' (一番小さい文字)をA_{i+1} - A_i個追加、A_i >= A_{i+1}ならS_iを切り詰めてインクリメント(最後の文字から1つ増やし、繰り上がりを左の文字へ伝播させていくこと)をしたものをS_{i+1}にすればよい。
愚直に実装するとO(A_i)時間、O(A_i)空間かかるが、A_iが大きい場合はほとんどの文字は'a'なので、'a'でない部分だけを管理すればよい。これはC++vector (RustのVec)などで実現できる。
計算量はO(N log N)。

実装上の注意点

  • 番兵としてS_0 = ""を用意すると楽
  • 最小化の問題なので、サンプルが合ったからといって信用しすぎず、挙動が合っているか(デバッグプリントなどで)調べる
  • 実装によってはk = 1のときがコーナーケースになるので気をつける

提出: Submission #3829390 - AtCoder Grand Contest 029 (Rust)

fn ok(a: &[i64], k: usize) -> bool {
    let mut cur: Vec<(_, usize)> = Vec::new();
    let mut curlen = 0;
    for &a in a {
        if curlen < a {
            curlen = a;
            continue;
        }
        // If k = 1, decrease the length and it's over.
        if k == 1 { return false; }
        curlen = a;
        // Truncate the managed string to a chars
        while let Some((x, elem)) = cur.pop() {
            if x >= a { continue; }
            cur.push((x, elem));
            break;
        }
        let mut carry = Vec::new();
        let mut pos = a - 1;
        // Increment
        loop {
            if let Some((x, elem)) = cur.pop() {
                if x == pos {
                    let c = (elem + 1) / k;
                    carry.push((x, (elem + 1) % k));
                    if c == 0 { break; }
                } else {
                    cur.push((x, elem));
                    carry.push((pos, 1));
                    break;
                }
            } else {
                carry.push((pos, 1));
                break;
            }
            if pos >= 1 { pos -= 1; }
            else { return false; }
        }
        for t in carry.into_iter().rev() { cur.push(t); }
    }
    true
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
        a: [i64; n],
    }
    let mut pass = n;
    let mut fail = 0;
    while pass - fail > 1 {
        let mid = (pass + fail) / 2;
        if ok(&a, mid) {
            pass = mid;
        } else {
            fail = mid;
        }
    }
    puts!("{}\n", pass);
}

まとめ

多倍長整数のインクリメントを実装しろ」をオブラートに包むとこういう問題になるのかな。

yukicoder No.753 最強王者決定戦

まあ基本だよね。
No.753 最強王者決定戦 - yukicoder

問題

16人の人間がいる。この16人を一列に並べて順番に組を作ってトーナメントを行ったとき、並べ方に応じて優勝者がただ一人決まる。全てのペアについて、戦ったときどちらが勝つかは決まっている。16!通りの並べ方全てについて優勝者を求めたとき、各人が何回優勝するかを計算せよ。

解法

DPで解く。DPテーブルは以下のようにする。
DP[i][S]: 集合Sの中の人が全員戦った時に人iが優勝する場合の数。ただし、組み方が同じで人の並べ方だけが違うものは同一視する。
(|S|は2べきのみを取り得ることに注意。)
DPの遷移はDP[i][S] <- DP[i][T] * DP[j][S - T] (iがjに勝つ場合)である。

計算量解析はかなり複雑。とりあえずΩ(2^n n^2)とO(4^n n^2)は確定だけどその間が詰められない。

実装上の注意点

  • 16!が64bit整数に収まるか注意する。今回はスターリングの公式から16! ~= (16/e)^16 < 8^16 = 2^48なので問題ない。
  • 同じトーナメントを何回も数えるので、最後に然るべき数を掛けるのを忘れないようにする。(今回は2^15)

提出: #305265 No.753 最強王者決定戦 - yukicoder (Rust)

/// Generates an Iterator over subsets of univ, in the descending order. 
/// Verified by: http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=3050308
struct SubsetIter { bits: Option<usize>, univ: usize }
impl Iterator for SubsetIter {
    type Item = usize;
    fn next(&mut self) -> Option<usize> {
        match self.bits {
            None => None,
            Some(bits) => {
                let ans = bits;
                self.bits =
                    if bits == 0 { None }
                else { Some((bits - 1) & self.univ) };
                Some(ans)
            }
        }
    }
}
fn subsets(univ: usize) -> SubsetIter {
    SubsetIter { bits: Some(univ), univ }
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    let n = 16;
    input! {
        a: [[i32; n]; n]
    }
    let mut a = a;
    for i in 0 .. n {
        for j in i + 1 .. n {
            a[j][i] = -a[i][j];
        }
    }
    let mut dp = vec![vec![0i64; 1 << n]; n];
    for i in 0 .. n {
        dp[i][1usize.wrapping_shl(i as u32)] = 1;
    }
    let mut s = 1;
    while s < n as u32 {
        s *= 2;
        for bits in 0 .. 1usize << n {
            if bits.count_ones() != s { continue; }
            for sub in subsets(bits) {
                if sub.count_ones() != s / 2 { continue; }
                let sub2 = bits - sub;
                for i in 0 .. n {
                    if (sub & 1 << i) == 0 { continue; }
                    for j in 0 .. n {
                        if (sub2 & 1 << j) == 0 { continue; }
                        if a[i][j] == 1 {
                            dp[i][bits] += dp[i][sub] * dp[j][sub2];
                        }
                    }
                }
            }
        }
    }
    for i in 0 .. n {
        puts!("{}\n", dp[i][(1 << n) - 1] << (n - 1));
    }
}

まとめ

なんで入力長を固定にしたり対称行列にしなかったりしたんだ…?

yukicoder No.770 Median Sequence

思いついた解法が嘘解法で、正しい解法を思いつくためにかなり多量の実験をした。
No.770 Median Sequence - yukicoder

問題

正の整数Nが与えられる。要素が1からNまでの長さNの整数列bに対して、以下の操作で数列aを作る。

  • a[1] = med(1, b[1], N)
  • a[2] = med(med(1, b[1], N - 1), b[2], N)
  • a[3] = med(med(med(1, b[1], N - 2), b[2], N - 1), b[3], N)
  • ...
  • a[i] = (m = 1; for(int j = 1; j <= i; ++j) m = med(m, b[i], N - i + j); を実行したときのmの値)

このとき、aとしてあり得るものの個数を数えよ。

解法

aとしてあり得るものは以下の操作で作れるものであり、およびそれに限る。
ここで、数列{c[1], ..., c[m]} が減少的であるとは、最後にc[i] < c[i+1]が成立する添字より最後にc[i] > c[i+1]が成立する添字が大きいことをいうことにする。
どちらかが存在しないときは、存在する方が大きいものとする。どちらも存在しないときはどちらも大きくないものとする。

  • a[1] は任意
  • 数列{a[1], ..., a[i - 1]}が減少的である場合、あるいはk = N - a[i - 1]としてi > k, a[i - k] = a[i - k + 1] = ... = a[i - 1]の両方が成り立つ場合、a[i] >= a[i - 1] - 1となるようにa[i]を選ぶ
  • それ以外の場合、a[i] >= a[i - 1]となるようにa[i]を選ぶ

これは以下のようなDPで計算できる。

DP[i][j]: 長さiで最後の要素がjであるような、減少的でない数列の個数
DP2[i][j]: 長さiで最後の要素がjであるような、減少的である数列の個数

実装上の注意点

  • 遷移が合っているかどうかを確認する
  • 答えが合わなかったら愚直解法と比較する

提出: #305095 No.770 Median Sequence - yukicoder (Rust)
#305093 No.770 Median Sequence - yukicoder (愚直解法のコードを含む)

const MOD: i64 = 1_000_000_007;
define_mod!(P, MOD);
type ModInt = mod_int::ModInt<P>;

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
    }
    let mut dp = vec![vec![ModInt::new(0); n]; n + 1];
    let mut dp2 = vec![vec![ModInt::new(0); n]; n + 1];
    dp[0][0] = ModInt::new(1);
    for i in 1 .. n + 1 {
        let mut cur = ModInt::new(0);
        for j in 0 .. n {
            cur += dp[i - 1][j];
            dp[i][j] += cur;
            cur += dp2[i - 1][j];
        }
        for j in 0 .. n - 1 {
            let turn = n - 1 - j;
            if i >= turn {
                dp2[i][j] += dp[i - turn][j + 1]
                    + dp2[i - 1][j + 1] + dp2[i - 1][j];
            }
        }
    }
    let mut tot = ModInt::new(0);
    for i in 0 .. n { tot += dp[n][i] + dp2[n][i]; }
    puts!("{}\n", tot);
}

まとめ

最初a[i]>=a[i-1]-1ならOKだと思ってサンプルが盛大に合わず苦しんだ。(そのような数列の個数はN=6のとき4488個ある。)