ARC086 F - Shift and Decrement

最初面倒 + 間違っている方針を考えていたが、某にもっと簡単な方法があると言われたので簡単な方法に気づくことができた。
F - Shift and Decrement

問題

N個の非負整数A_1, ..., A_Nが黒板に書かれている。以下の操作をK回まで行う。

  • 以下の2個から一つ選ぶ。
    • 操作A: 黒板の全ての整数xについて、それをfloor(x/2)で置き換える。
    • 操作B: 黒板の全ての整数xについて、それをx - 1で置き換える。これは全ての整数が1以上であるときにのみ行える。

最終的に黒板に書かれている非負整数の組み合わせとしてあり得るものをmod (10^9 + 7)で出力せよ。

解法

まず、どちらの操作も広義単調増加な変換しかしないので、操作によって整数の大きさの順番が変わることはない。よって、Aは多重集合ではなく数列として扱うことができる。Aは昇順にソートされていると仮定してよい。
どんな状態へも、[(操作B1回以下) -> (操作A)]の有限回の繰り返し -> 操作Bの有限回の繰り返しで到達することができる。最後の操作Aの前に操作Bを2回行う場合、それを直後の操作Aの後の操作B1回に置き換えられるためである。最適(操作回数最小)なものがこの形に限ることもほぼ明らか。
1を長さNのすべて1からなる数列とし、減算や乗算を項ごとに定めることにする。[(操作B1回以下) -> (操作A)]の有限回の繰り返しでAからvに到達したとすると、そこからはv, v - 1, v - 21, ..., v -k1に行くことができる。
逆に、到達可能な状態は全てこのように表せるので、上の(v, k)を全列挙すれば良い。これにはメモ化再帰が使える。
さて、このように(v, k)が列挙されたとき、それをまとめることを考えよう。vの各要素からvの最小値を引いてできる配列をd(v)としたとき、v1 - s1 = v2 - t1が成立するのはd(v1) = d(v2), v1[1] - s = v2[1] - tのとき、及びそのときに限る。(仮定からv1もv2もソートされていることに注意。)
以上から、d(v)が同じ状態をまとめると、それぞれについてあり得る最終状態はd(v) + u1 (uは然るべき区間のunionの上を動く)であることが分かる。各d(v)について、uが動く集合(下記コードのrangesのunion)の大きさを求めればよい。これは区間の出現・消失イベントの起きる時刻でソートすればできる。
計算量はどうなるだろうか? 最初のメモ化再帰で訪れる頂点の個数を見積もる。操作Aを使う個数を固定し、最後に操作Aを使った状態だけを考えると、現れる数列はA/2^b - 1からA/2^bの間に収まるので、高々N通りである。探索順を工夫すれば、各状態を2回以上探索しないようにできるため、探索すべき状態の個数はO(N * log max A)程度である。1状態ごとにかかる時間はO(N)なので、全体の計算量はO(N^2 * log max A)である。

登場する典型

  • 最適戦略についての考察
  • メモ化再帰
  • 区間のunionの大きさ

実装上の注意点

  • traverseの順番に気をつける
    • kが大きい順にtraverseすることで、ある状態に2回目以降訪れた場合にそこから先の探索が起きないことが保証できる(ダイクストラ法に近い)

提出: Submission #3974898 - AtCoder Regular Contest 086 (Rust)

fn chmax<T: Ord>(a: &mut T, b: T) -> bool {
    if *a >= b { return false; }
    *a = b;
    true
}

fn dfs(a: Vec<i64>, k: i64, memo: &mut HashMap<Vec<i64>, i64>) {
    {
        let entry = memo.entry(a.clone()).or_insert(-1);
        if !chmax(entry, k) { return }
    }
    if k == 0 { return };
    let half = a.iter().map(|&x| x / 2).collect();
    dfs(half, k - 1, memo);
    if k > 1 && a[0] > 0 {
        let half = a.iter().map(|&x| (x - 1) / 2).collect();
        dfs(half, k - 2, memo);
    }
}

const MOD: i64 = 1_000_000_007;

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        k: i64,
        a: [i64; n],
    }

    let mut a = a;
    a.sort();

    let mut memo = HashMap::new();
    dfs(a, k, &mut memo);
    // compress
    let mut comp = HashMap::new();
    for (a, k) in memo {
        let mi = a[0];
        let mut diff = vec![0; n - 1];
        for i in 0 .. n - 1 { diff[i] = a[i + 1] - mi; }
        let range = (max(0, mi - k), mi);
        let entry = comp.entry(diff).or_insert(Vec::new());
        entry.push(range);
    }
    let mut tot = 0;
    for (_, ranges) in comp {
        let mut events = Vec::new();
        for (x, y) in ranges {
            events.push((x, 1));
            events.push((y + 1, -1));
        }
        events.sort();
        let mut cur = 0;
        let mut last = 0;
        for (t, delta) in events {
            if cur > 0 {
                tot += t - last;
                tot %= MOD;
            }
            last = t;
            cur += delta;
        }
    }
    puts!("{}\n", tot);
}

まとめ

本番誰も解けなかったのが意外である。ちなみに私は時間ギリギリにEを通して狂喜してた。