AOJ 3165 Triangle Addition (HUPC 2020 day1-B)

いもす法

• 問題へのリンク
• editorial

問題概要

要素からなる数列がある。初期状態では全要素の値が 0 である。以下の 回の操作を行なって得られる数列を出力せよ。

• 各クエリでは整数 が与えられる。区間 に対して、初項 1、交差 1 の等差数列を加算する

制約

• 

考えたこと

(0, 0, 1, 2, 3, 4, 5, 0, 0) といった値を加算することになる。この手のいもす法っぽい問題では、「階差をとる」という操作をやってみると見えてくる。まず 1 回階差をとると

(0, 0, 1, 1, 1, 1, 1, -5, 0)

となる。さらに階差をとると

(0, 0, 1, 0, 0, 0, 0, -6, 5)

となる。以上から

• 各クエリに対して、a[ L ] += 1、a[ L + K ] -= K + 1、a[ L + K + 1 ] += K とする
• 最後に 2 回累積和をとる

とすればよいことがわかる。計算量は 。

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

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<long long> a(N, 0);
    for (int q = 0; q < Q; ++q) {
        int L, K;
        cin >> L >> K;
        --L;
        a[L]++;
        if (L + K < N) a[L + K] -= K + 1;
        if (L + K + 1 < N) a[L + K + 1] += K;
    }
    for (int i = 0; i + 1 < N; ++i) a[i + 1] += a[i];
    for (int i = 0; i + 1 < N; ++i) a[i + 1] += a[i];
    for (int i = 0; i < N; ++i) {
        if (i) cout << " ";
        cout << a[i];
    }
    cout << endl;
}

別解

遅延評価セグメントツリーで殴ることもできる。等差数列の加算は「初項 (第 0 項目換算)」と「交差」を管理することで実現できる。

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

// Segment Tree
template<class Monoid, class Action> struct SegTree {
    using FuncMonoid = function< Monoid(Monoid, Monoid) >;
    using FuncAction = function< void(Monoid&, Action) >;
    using FuncComposition = function< void(Action&, Action) >;
    FuncMonoid FM;
    FuncAction FA;
    FuncComposition FC;
    Monoid IDENTITY_MONOID;
    Action IDENTITY_LAZY;
    int SIZE, HEIGHT;
    vector<Monoid> dat;
    vector<Action> lazy;
    
    SegTree() {}
    SegTree(int n, const FuncMonoid fm, const FuncAction fa,
            const FuncComposition fc,
            const Monoid &identity_monoid, const Action &identity_lazy)
    : FM(fm), FA(fa), FC(fc), 
      IDENTITY_MONOID(identity_monoid), IDENTITY_LAZY(identity_lazy) {
        SIZE = 1, HEIGHT = 0;
        while (SIZE < n) SIZE <<= 1, ++HEIGHT;
        dat.assign(SIZE * 2, IDENTITY_MONOID);
        lazy.assign(SIZE * 2, IDENTITY_LAZY);
    }
    void init(int n, const FuncMonoid fm, const FuncAction fa,
              const FuncComposition fc,
              const Monoid &identity_monoid, const Action &identity_lazy) {
        FM = fm, FA = fa, FC = fc;
        IDENTITY_MONOID = identity_monoid, IDENTITY_LAZY = identity_lazy;
        SIZE = 1, HEIGHT = 0;
        while (SIZE < n) SIZE <<= 1, ++HEIGHT;
        dat.assign(SIZE * 2, IDENTITY_MONOID);
        lazy.assign(SIZE * 2, IDENTITY_LAZY);
    }
    
    // set, a is 0-indexed
    void set(int a, const Monoid &v) { dat[a + SIZE] = v; }
    void build() {
        for (int k = SIZE - 1; k > 0; --k)
            dat[k] = FM(dat[k*2], dat[k*2+1]);
    }
    
    // update [a, b)
    inline void evaluate(int k) {
        if (lazy[k] == IDENTITY_LAZY) return;
        if (k < SIZE) FC(lazy[k*2], lazy[k]), FC(lazy[k*2+1], lazy[k]);
        FA(dat[k], lazy[k]);
        lazy[k] = IDENTITY_LAZY;
    }
    inline void update(int a, int b, const Action &v, int k, int l, int r) {
        evaluate(k);
        if (a <= l && r <= b) FC(lazy[k], v), evaluate(k);
        else if (a < r && l < b) {
            update(a, b, v, k*2, l, (l+r)>>1);
            update(a, b, v, k*2+1, (l+r)>>1, r);
            dat[k] = FM(dat[k*2], dat[k*2+1]);
        }
    }
    inline void update(int a, int b, const Action &v) { 
        update(a, b, v, 1, 0, SIZE);
    }
    
    // get [a, b)
    inline Monoid get(int a, int b, int k, int l, int r) {
        evaluate(k);
        if (a <= l && r <= b)
            return dat[k];
        else if (a < r && l < b)
            return FM(get(a, b, k*2, l, (l+r)>>1), 
                      get(a, b, k*2+1, (l+r)>>1, r));
        else
            return IDENTITY_MONOID;
    }
    inline Monoid get(int a, int b) { 
        return get(a, b, 1, 0, SIZE);
    }
    inline Monoid operator [] (int a) {
        return get(a, a + 1);
    }
    
    // debug
    void print() {
        for (int i = 0; i < SIZE; ++i) {
            if (i) cout << ",";
            cout << get(i, i+1);
        }
        cout << endl;
    }
};

using pll = pair<long long, long long>;
int main() {
    int N, Q;
    cin >> N >> Q;

    auto fm = [&](pll a, pll b) {
        return pll(a.first + b.first, a.second + b.second);
    };
    auto fa = [&](pll &a, pll d) {
        a.first += d.first, a.second += d.second;
    };
    auto fc = [&](pll &d, pll e) {
        d.first += e.first, d.second += e.second;
    };
    pll identity_monoid = pll(0, 0);
    pll identity_lazy = pll(0, 0);
    SegTree<pll, pll> seg(N, fm, fa, fc, identity_monoid, identity_lazy);

    vector<long long> a(N, 0);
    for (int q = 0; q < Q; ++q) {
        int L, K;
        cin >> L >> K;
        --L;
        int R = min(L + K, N);
        int syoko = 1 - L;
        seg.update(L, R, pll(syoko, 1));
    }
    for (int i = 0; i < N; ++i) {
        if (i) cout << " ";
        auto p = seg.get(i, i+1);
        cout << p.first + p.second * i;
    }
    cout << endl;
}