➕

[MojaCoder] お茶コン02 - 04「区間和の和」Tester 解法

2023/11/13に公開

競技プログラミング

mojacoder

tech

お茶コン02 の Tester です。
04. 区間和の和について、「別解」というほど異なりはしませんが、ちょっと違う実装方針だったのでメモです。

次のように $\oplus\, \colon\, \N^3 \times \N^3 \to \N^3$ と $\hat 0 \in \N^3$ とを定めます：

$(V_0, S_0, I_0) \oplus (V_1, S_1, I_1) = (V_0 + V_1 + S_0 \times I_1,\, S_0 + S_1,\, I_0 + I_1)$
$\hat 0 = (0, 0, 0)$

このとき， $\displaystyle V_{L, R} \coloneqq \!\!\!\! \sum_{L \leq j \leq i \leq R} \!\!\! A_j$ に対して以下が成り立ちます：

$\displaystyle (V_{L, R}, S_{L, R}, I_{L, R}) = \bigoplus_{L \leq i \leq R} \; (A_i, A_i, 1)$
- ただし，以下を満たす：
  - $\displaystyle S_{L, R} = \sum_{L \leq i \leq R} \!\! A_i$
  - $I_{L, R} = R - L + 1$

組 $(\N^3, \oplus, \hat 0)$ は monoid の公理を満たしますから Segment Tree のようなデータ構造に乗ります．

実装例

C++

#include <iostream>
#include <atcoder/modint>
#include <atcoder/segtree>

using mint = atcoder::modint998244353;

// triple (val, sum, len)
struct S {
    mint val, sum;
    int len = 0;
    S() = default;
    S(mint v, mint s, int l) : val(v), sum(s), len(l) {}
    S(mint v) : val(v), sum(v), len(1) {}
    operator mint() const { return this->val; }
};

S op(S p0, S p1) {
    auto [ v0, s0, l0 ] = p0;
    auto [ v1, s1, l1 ] = p1;
    return {
        v0 + v1 + s0 * l1,
        s0 + s1,
        l0 + l1,
    };
}

S e() { return { 0, 0, 0 }; }

signed main() {
    int n; std::cin >> n;

    atcoder::segtree<S,op,e> data(
        ({
            std::vector<S> a(n);
            for(int i=0; i<n; ++i) {
                int v; std::cin >> v;
                a[i] = mint{ v };
            }
            a;
        })
    );

    int q; std::cin >> q;
    while(q--) {
        int t; std::cin >> t;
        if(t == 1) {
            int l, r; std::cin >> l >> r; --l;
            std::cout << mint{ data.prod(l, r) }.val() << "\n";
        }
        if(t == 2) {
            int k, x; std::cin >> k >> x; --k;
            data.set(k, mint{ x });
        }
    }
}

__attribute__((constructor)) void fast_io() { std::ios::sync_with_stdio(false), std::cin.tie(nullptr); }

実装例

Discussion