ABC 340

Table of Contents

E - Mancala 2

https://atcoder.jp/contests/abc340/tasks/abc340_e

区間加算、一点取得の問題。ac-library を使うと lazy segment tree が必要になってセットアップが大変だが自分でセグメントツリーを実装すると意外と簡単。

かつっぱ氏の木マスター養成講座がわかりやすい ref: https://youtu.be/g30mEb4jE2g?si=g-oLairJqPWzWHn6

struct SegmentTree {
    vll seg;
    int len;

    SegmentTree(int n) {
        rep(i, 30) {
            if (n <= (1 << i)) {
                len = 1 << i;
                break;
            }
        }
        seg.resize(len * 2, 0);
    }

    ll get(int p) {
        p += len;
        ll ans = seg[p];
        while (p / 2) {
            p /= 2;
            ans += seg[p];
        }

        return ans;
    }

    void add(int l, int r, ll x) {
        l += len, r += len;
        while (l < r) {
            if (l & 1) {
                seg[l] += x;
                l++;
            }
            l /= 2;
            if (r & 1) {
                seg[r - 1] += x;
                r--;
            }
            r /= 2;
        }
    }
};

void solve() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    ll n, m;
    cin >> n >> m;

    SegmentTree seg(n);
    rep(i, n) {
        ll a;
        cin >> a;
        seg.add(i, i + 1, a);
    }

    rep(i, m) {
        int b;
        cin >> b;
        ll x = seg.get(b);
        seg.add(b, b + 1, -x);

        b++;
        ll q = x / n;
        ll r = x % n;
        seg.add(0, n, q);
        if (b + r >= n) {
            seg.add(b, n, 1);
            seg.add(0, b + r - n, 1);
        } else {
            seg.add(b, b + r, 1);
        }
    }

    rep(i, n) {
        if (i != 0)
            cout << ' ';
        cout << seg.get(i);
    }
    cout << endl;
}