AtCoder Grand Contest 006 D - Median Pyramid Hard

問題記録 AtCoder AGC AGC #006 二分探索

http://agc006.contest.atcoder.jp/tasks/agc006_d

問題概要

N段のピラミッドがあり、i段目は(2i-1)個のブロックからなる。最下段の(2N-1)個のブロックには(1, 2, …, 2N-1)を並べ替えた数字がそれぞれ書き込まれている。最下段以外のi段目のブロックのj番目の数字は、(i+1)段目のj番目, (j+1)番目, (j+2)番目のブロックの中央値となる。いま最下段に書き込まれた数字の列が与えられるので、最上段の1個のブロックに書き込まれることになる数を答えよ。

解法

2分探索。中では以下のような計算を回す。

数xをひとつ定め、最下段のすべての数字を

  • x以上ならば1
  • x未満ならば0

と書き換える。この0/1数列で元と同様の操作を行った結果として最上段が1になるならば答えはx以上の範囲にあり、0になるならば答えはx未満の範囲にあることになる。これで2分探索の範囲を絞っていくことができる。そしてこの判定は以下のように行うことで十分高速にできる。

  • 最下段の真ん中3要素が 111, 110, 011 のとき → 最上段は1になる
  • 最下段の真ん中3要素が 000, 001, 100 のとき → 最上段は0になる
  • 最下段の真ん中3要素が 101, 010 のとき → 真ん中から横に見ていって、途中で同じ数が連続した場合にはその数が最上段になる。仮にそのような数がなければ（＝数列が010101…の繰り返しの場合）、ひとつ上の段に行くたび真ん中の数が反転することになる。

2分探索の中ではO(N)かかるので全体でO(NlogN)になって十分高速に求まる。

感想

この問題すごすぎないですか？リアルタイムで参加してて解説読んでめちゃくちゃ感動した記憶があります

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
void main() {
    auto N = readln.chomp.to!int;
    auto M = 2 * N - 1;
    auto A = readln.split.map!(to!int).array;
 
    int hi = M + 1;
    int lo = 0;
    while (hi - lo > 1) {
        int mid = (hi + lo) / 2;
        auto B = A.map!(a => a >= mid).array;
        int a = B[N-2], b = B[N-1], c = B[N];
        if ((a & b) || (b & c)) {
            lo = mid;
        } else if ((!a & !b) || (!b & !c)) {
            hi = mid;
        } else {
            bool ok = false;
            for (int l = N - 3, r = N + 1; l >= 0; l--, r++) {
                if (B[l] == B[l+1]) {
                    if (B[l]) lo = mid;
                    else hi = mid;
                    ok = true;
                    break;
                } else if (B[r] == B[r-1]) {
                    if (B[r]) lo = mid;
                    else hi = mid;
                    ok = true;
                    break;
                }
            }
            if (!ok) {
                bool hoge = B[N-1];
                if (M / 2 % 2) hoge ^= 1;
                if (hoge) lo = mid;
                else hi = mid;
            }
        }
    }
 
    lo.writeln;
}