square869120Contest #2 H - Counting 1's

問題

H: Counting 1's - square869120Contest #2 | AtCoder

N個の0が並んでいる。

以下のクエリを処理せよ。

• クエリ1: 区間[l, r)のビットを反転させる。
• クエリ2: 区間[l, r)内の1の数を答える。

解法

遅延セグメント木。区間の問題なのでセグ木という発想は出てくるんですが、クエリ1の反転処理を1ビットずつやっていたのでは時間がかかります。そこで遅延評価です。各ノードに「後で反転させる予定があるか」を記録しておきます。反転処理は、自分を反転させたら子に「反転させる予定」を押し下げるということを通ったノードだけします。つまり毎回すべてのビットを反転させる必要はなくて、知りたくなったら処理しましょうという感じですね。 （よくわかってないので変なこと言ってたら教えて）

あとC++でデフォルト引数が使えるのを初めて知ったので使ってみました。

#include <bits/stdc++.h>
using namespace std;
#define int long long   // <-----!!!!!!!!!!!!!!!!!!!

#define rep(i,n) for (int i=0;i<(n);i++)
#define rep2(i,a,b) for (int i=(a);i<(b);i++)
#define rrep(i,n) for (int i=(n)-1;i>=0;i--)
#define rrep2(i,a,b) for (int i=(b)-1;i>=(a);i--)
#define all(a) (a).begin(),(a).end()

typedef long long ll;
typedef pair<int, int> P;

const int MAX_N = 1 << 17;
int seg[MAX_N * 2];     // 区間内の1の個数
int lazy[MAX_N * 2];    // 後でノードを反転させるつもりがあるか

// 遅延評価のアレ
void push(int k, int l, int r) {
    // このノードを反転させるつもりがなかったなら何もしない
    if (lazy[k] == 0) return;

    // 以下、このノードを反転させる場合
    // ノードが表す区間内の1の数が反転
    seg[k] = (r - l) - seg[k];

    // ノードkが子を持つならば、子を後で反転させるつもりにしておく
    if (r - l > 1) {
        lazy[k * 2 + 1] ^= 1;
        lazy[k * 2 + 2] ^= 1;
    }

    // ノードkはもう反転させたので0にしておく
    lazy[k] = 0;
}

// 区間内のビットを反転させる
void inv(int a, int b, int k = 0, int l = 0, int r = MAX_N) {
    push(k, l, r);
    // ノードが範囲外なら何もしない
    if (r <= a || b <= l) {
        return;
    }
    // ノードが完全に区間に含まれていれば、反転回数が反転
    if (a <= l && r <= b) {
        lazy[k] ^= 1;
        push(k, l, r);
    }
    // 子ノードを反転させ（↓）、自身を子ノードの和に更新（↑）
    else {
        inv(a, b, k * 2 + 1, l, (l + r) / 2);
        inv(a, b, k * 2 + 2, (l + r) / 2, r);
        seg[k] = seg[k * 2 + 1] + seg[k * 2 + 2];
    }
}


int sum(int a, int b, int k = 0, int l = 0, int r = MAX_N) {
    push(k, l, r);
    // ノードが範囲外なら0を返す
    if (r <= a || b <= l) {
        return 0;
    }
    // ノードが完全に区間に含まれていれば、このノードの値を返す
    else if (a <= l && r <= b) {
        return seg[k];
    }
    // 上記以外の場合、左右の子ノードの和を返す
    else {
        int res = 0;
        res += sum(a, b, k * 2 + 1, l, (l + r) / 2);
        res += sum(a, b, k * 2 + 2, (l + r) / 2, r);
        return res;
    }
}

signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);

    int N, Q;
    cin >> N >> Q;
    rep(i, Q) {
        int q, l, r;
        cin >> q >> l >> r;
        if (q == 1) {
            inv(l, r);
        } else {
            cout << sum(l, r) << endl;
        }
    }
}

セグ木が少し面白いと思えるようになってきた。