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

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

AtCoder ABC 129 E - Sum Equals Xor (500 点)

まさにこの考え方で解ける問題

drken1215.hatenablog.com

今回の問題も桁 DP でもよいのだが、実は桁 DP しなくても解ける。

問題へのリンク

問題概要

正整数  L が二進数表記で与えられます。 以下の条件を満たす非負整数  a, b の組がいくつ存在するか求めてください (1000000007 で割ったあまり)

  •  a+b \le L
  •  a+b = a XOR  b

制約

  •  L の桁数は  10^{5} 以下

「a + b = a XOR b」の言い換え

 L 以下」という条件がある時点で、桁DP か、それに類する発想を使いそうな問題に思える。

それを念頭に置きつつ、まずは「a + b = a XOR b」という条件をわかりやすく言い換えることから始めてみる。このように「わかりにくい条件をわかりやすく言い換える」というのは AtCoder では定番なイメージがある。

さて、a と b をともに二進法表記したとき

a = 1010
b = 0110

という風だったとき、下から二番目の位に着目すると、a も b も 1 になっている。このとき

  • a + b ではくり上がりが起こる
  • a XOR b ではくり上がりが起こらず、相殺して 0 になる

ということが読み取れる。反対に (a の i 桁目, b の i 桁目) = (1, 0), (0, 1), (0, 0) の場合にはいずれも i 桁目については a+b と a XOR b とは等しくて、しかも他の桁への影響もない。

よって、


a + b >= a XOR b が常に成立していて、等号成立条件は「どの桁についても a と b がともに 1 になることはないこと」である


ということが言える。

元の問題へ

よって問題は

  • a + b <= L
  • どの桁をみても a と b の少なくとも片方は 0 である

というものを数え上げる問題となった。editorial では桁 DP でやっているが、ここでは桁 DP の思想を引き継ぎつつ表立って DP せずにやってみる。

まず、例えば L = 11001010 という整数だったとき、L 以下の整数というのは以下のように場合分けできる。ここで「...」の部分はなんでもよいことを表す。

  • 0...
  • 10...
  • 11000...
  • 1100000...
  • 11001010

それぞれについて、

  • 「...」の部分については、その部分の桁数を d とすると、各桁について (a, b) = (0, 0), (0, 1), (1, 0) の 3 通りがあるので、 3^{d} 通り
  • それ以外の部分については、1 の個数を e とすると、1 が立っている各桁について (a, b) = (0, 1), (1, 0) の 2 通りがあるので、 2^{e} 通り。

ということなので、まとめると  3^{d} \times 2^{e} 通りとなる。これを合算すれば OK。

#include <iostream>
#include <string>
using namespace std;


// 1000000007 で割ったあまりを扱う構造体
template<int MOD> struct Fp {
    long long val;
    constexpr Fp(long long v = 0) noexcept : val(v % MOD) {
        if (val < 0) v += MOD;
    }
    constexpr int getmod() { return MOD; }
    constexpr Fp operator - () const noexcept {
        return val ? MOD - val : 0;
    }
    constexpr Fp operator + (const Fp& r) const noexcept { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp& r) const noexcept { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp& r) const noexcept { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp& r) const noexcept { return Fp(*this) /= r; }
    constexpr Fp& operator += (const Fp& r) noexcept {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp& r) noexcept {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp& r) noexcept {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp& r) noexcept {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b; swap(a, b);
            u -= t * v; swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr bool operator == (const Fp& r) const noexcept {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp& r) const noexcept {
        return this->val != r.val;
    }
    friend constexpr ostream& operator << (ostream &os, const Fp<MOD>& x) noexcept {
        return os << x.val;
    }
    friend constexpr istream& operator >> (istream &is, Fp<MOD>& x) noexcept {
        return is >> x.val;
    }
    friend constexpr Fp<MOD> modpow(const Fp<MOD> &a, long long n) noexcept {
        if (n == 0) return 1;
        auto t = modpow(a, n / 2);
        t = t * t;
        if (n & 1) t = t * a;
        return t;
    }
};
const int MOD = 1000000007;
using mint = Fp<MOD>;


mint solve(const string &L) {
    long long num = 0;
    mint res = 0;
    for (int i = 0; i < L.size(); ++i) {
        int len = (int)L.size() - i - 1;
        if (L[i] == '1') {
            res += modpow(mint(2), num) * modpow(mint(3), len);
            ++num;
        }
    }
    res += modpow(mint(2), num);
    return res;
}

int main() {
    string L; cin >> L;
    cout << solve(L) << endl;
}