けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 301 D - Bitmask (緑色, 400 点)

AtCoder AtCoder400点 ABC-D どの場所で初めてsmallerになるかを考える 辞書順 DP状態:smaller(桁DP) Greedy 0と1の問題 ワイルドカード問題 ある値を固定して考える Greedy:ある量を決めると残りが決まっていく 緑色diff Greedy:辞書順最小を求める テク:「最大値がK以下」⇔「すべてがK以下」 Yes/No判定問題 最大スコア 条件の言い換え 制約条件:ある値<=K DP状態:どの桁まで見たか(広義の桁DP) 最適化問題

基本的に Greedy にやればよさそうなのに、意外とやりづらい問題。基本に忠実にやれば解ける！

問題概要

文字 '0', '1', '?' のみからなる文字列 S が与えられる。

S 中の各 '?' をそれぞれ '0' または '1' に置き換えて得られる文字列を二進法表記の整数値とみなすことにする。

それが整数 N 以下であるようにするとき、得られる整数値の最大値を求めよ。ただし、どのように '?' を置き換えても N 以下にならない場合は -1 を出力せよ。

制約

  • 1 \le |S| \le 60
  • 1 \le N \le 10^{18}

考えたこと：嫌な場合

基本的には、 N を超えない範囲で、? を 1 に置き換えていきたい。少し嫌なケースが次のような場合だ ( N を二進法表記している)。

  • S = "??1?"
  • N = "1101"

このとき、先頭の桁から順番に埋めようとすると、 N の先頭の 2 文字が 1 であることから、 S の先頭 2 個の ? を 1 で埋めたくなる。しかしこのとき、

  • S = "111?"
  • N = "1101"

となる。この場合、 S が N を超えてしまうのだ。こういうことがあるから、何も考えずに Greedy といわけにはいかない。ここで Greedy 戦略の基本に立ち返ろう。

解法 1：それでも、先頭の桁から貪欲に見る

上記のような場合があっても、なお、先頭から Greedy に攻めていく戦法は有効だ。次のように考えれば良い。

  • 先頭の桁から順に見ていき
    • S のその桁が '?' 以外のとき:
      • その文字にする
    • S のその桁の '?' を 1 にすることが可能であるとき:
      • S のその桁の値を 1 にする
    • 可能でないとき:
      • S のその桁の値を 0 にする
  • こうして得られた値が N より大きいときはそもそもダメ

ここで問題になるのは、「 S のその桁の '?' を 1 にすることが可能かどうか」をいかにして判定するかだ。

これは実は簡単だ。

S のその桁の '?' を 1 にした上で理論上達成できる S の値の最小値が N 以下であるかどうかを判定すればよいのだ。具体的には、 S の残りの ? をすべて 0 にして N 以下になるかどうかを調べればよい。

コード

以上のロジックを実装すると、次のように書ける。この実装の計算量は O(|S|^{2}) となる。

なお、もっと簡単に O(|S|) の計算量のコードも書ける (→ evima さんの解説)

#include <bits/stdc++.h>
using namespace std;

int main() {
    string S;
    long long N;
    cin >> S >> N;
    
    // 先頭の桁から見ていく
    long long res = 0;
    for (int i = 0; i < S.size(); ++i) {
        if (S[i] == '1') res += 1LL<<(S.size()-1-i);
        else if (S[i] == '?') {
            // S のこの桁を 1 にした上で、残りの ? をすべて 0 にした値を求める
            long long smin = res + (1LL<<(S.size()-1-i));
            for (int j = i+1; j < S.size(); ++j) {
                if (S[j] == '1') smin += 1LL<<(S.size()-1-j);
            }
            
            // その値が T 以下なら 1 にする
            if (smin <= N) res += 1LL<<(S.size()-1-i);
        }
    }
    cout << (res <= N ? res : -1) << endl;
}

解法 2：「 N 以下」という条件をうまく扱う

もう一つ、重要な考え方に基づく解法をやってみる。こっちは、より公式解説に近い解法になる。

一般に、整数値 A が整数値 N 以下であるとは、次のことを意味する

ある整数 d が存在して

  • A と N の先頭から d 桁は等しい
  • その次の桁については、 A の値は N の値より小さい
  • それ以降の桁については、 A の値はなんでもよい

なお、この考え方は「桁 DP」をやる上でも、中心的な考え方となるものでもある。

この考え方に基づいて、次の解法が考えられる。

  • d = 0, 1, \dots, |S| として、
    • S, N の先頭から d 桁が等しくなるようにする (できない場合はスキップ)
    • その次の桁では S の値が 0、 N の値が 1 になるようにする (できない場合はスキップ)
    • 残りの S の ? はすべて 1 にする

各 d に対する解の最大値をとる

計算量は O(|S|^{2}) となる。

#include <bits/stdc++.h>
using namespace std;

int main() {
    string S;
    long long N;
    cin >> S >> N;
    
    // N の桁数がそもそも S より大きい場合: S の ? をすべて 1 に
    if (N >= (1LL<<S.size())) {
        long long res = 0;
        for (int i = 0; i < S.size(); ++i) {
            if (S[i] != '0') res += 1LL<<(S.size()-1-i);
        }
        cout << res << endl;
        return 0;
    }
    
    // 先頭から d 桁を一致させる場合を考える
    long long res = -1;
    for (int d = 0; d <= S.size(); ++d) {
        long long tmp = 0;
        
        // 先頭から d 桁を一致させられるかどうか
        bool ok = true;
        for (int i = 0; i < d; ++i) {
            if (N & (1LL<<(S.size()-1-i))) {
                tmp += 1LL<<(S.size()-1-i);
                if (S[i] == '0') ok = false;
            } else {
                if (S[i] == '1') ok = false;
            }
        }
        if (!ok) continue;

        // その次の桁で S を 0、N が 1 になる場合のみ OK
        if (d < S.size() && !(N & (1LL<<(S.size()-1-d)))) continue;
        if (d < S.size() && S[d] == '1') continue;
        
        // S の残りの ? を 1 にする
        for (int i = d+1; i < S.size(); ++i) {
            if (S[i] != '0') tmp += 1LL<<(S.size()-1-i);
        }
        res = max(res, tmp);
    }
    cout << res << endl;
}