DDCC 2018 D - DISCO! - すいバカ日誌

問題文

https://atcoder.jp/contests/ddcc2019-final/tasks/ddcc2019_final_d

解法

区間同士でうまくマージできるような形であればセグ木に乗せることができるので、とりあえずその方針で考えてみる。
2つの区間 A, B があって、それらをマージしたときに "DISCO" がいくつ現れるかを考えてみると

A から "", B から "DISCO"
A から "D", B から "ISCO"
A から "DI", B から "SCO"
A から "DIS, B から "CO"
A から "DISC", B から "O"
A から "DISCO", B から ""

の5通りある。
よって、それぞれの区間に、対応する文字列が何通りあるか持てば良いことがわかる。
この議論は、当然 "ISCO" や "SCO" などを数えるときにも必要であるが、
冷静に考えると、文字列は "DISCO" の連続部分文字列であり、先程列挙したように、文字列の構成の仕方に特徴がある。
これに気がつけば
A[i][j] := "DISCO" の [i, j) が完成している場合の数
と定義した行列積が、区間のマージに対応することがわかる。

たとえば、s[i] が 'I' であるなら、A[1][2] = 1 とすればよい（ただし対角成分は常に1）。

以上で、計算量は O(6^3 * Q * logN) となる。
TL が 13sec なので普通に書けば通るだろう。

ソースコード

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

using u32 = uint32_t;
using matrix = array<array<u32, 6>, 6>;

matrix make_zmat() {
    matrix res;
    for(int i = 0; i < 6; ++i) {
        for(int j = 0; j < 6; ++j) {
            res[i][j] = 0;
        }
    }
    return res;
}

matrix make_id_mat() {
    auto res = make_zmat();
    for(int i = 0; i < 6; ++i) res[i][i] = 1;
    return res;
}

matrix operator*(matrix const& a, matrix const& b) {
    auto res = make_zmat();
    for(int k = 0; k < 6; ++k) {
        for(int i = 0; i < 6; ++i) {
            for(int j = 0; j < 6; ++j) {
                res[i][j] += a[i][k] * b[k][j];
            }
        }
    }
    return res;
}

template<typename Monoid>
class segment_tree {
    using T = typename Monoid::type;

public:
    segment_tree(std::vector<T> const& init)
        : sz(init.size()), n(expand(init.size()))
    {
        dat.assign(n*2, Monoid::id());
        std::copy(begin(init), end(init), begin(dat) + n);
        for(int i = n - 1; i >= 0; --i) {
            dat[i] = Monoid::op(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    segment_tree(int const n_, T const& init = Monoid::id())
        : segment_tree(std::vector<T>(n_, init))
    {}

    void update(int p, T val) {
        assert(0 <= p && p < sz);
        dat[p += n] = val;
        while(p /= 2) {
            dat[p] = Monoid::op(dat[p * 2], dat[p * 2 + 1]);
        }
    }

    // [l, r)
    T query(int l, int r) const {
        assert(0 <= l && l <= r && r <= sz);
        l += n, r += n;
        T res1 = Monoid::id(), res2 = Monoid::id();
        while(l != r) {
            if(l & 1) res1 = Monoid::op(res1, dat[l++]);
            if(r & 1) res2 = Monoid::op(dat[--r], res2);
            l /= 2, r /= 2;
        }
        return Monoid::op(res1, res2);
    }

private:
    int expand(int n_) const {
        assert(n_ >= 1);
        return n_ == 1 ? n_ : expand((n_ + 1) / 2) * 2;
    }

private:
    const int sz, n;
    std::vector<T> dat;
};

struct DISCO {
    using type = matrix;

    static type id() {
        return make_id_mat();
    }
    static type op(type const& l, type const& r) {
        return l * r;
    }
};

int main() {
    string s; cin >> s;
    const int n = s.size();

    segment_tree<DISCO> seg(n);
    const string disco = "DISCO";
    for(int i = 0; i < n; ++i) {
        auto a = make_id_mat();
        const int idx = find(begin(disco), end(disco), s[i]) - begin(disco);
        a[idx][idx + 1] = 1;
        seg.update(i, a);
    }

    int q; cin >> q;
    while(q--) {
        int l, r; cin >> l >> r;
        l--;
        cout << seg.query(l, r)[0][5] << endl;
    }
}

感想

なぜか本番で解けなかったやつ。大反省。