Skip to content

Instantly share code, notes, and snippets.

@KirillLykov
Last active October 23, 2020 09:06
Show Gist options
  • Select an option

  • Save KirillLykov/940f046ac6347e60bfeda7773cdd613c to your computer and use it in GitHub Desktop.

Select an option

Save KirillLykov/940f046ac6347e60bfeda7773cdd613c to your computer and use it in GitHub Desktop.
Addition on a range using prefix sums and similar

Problem 1: Add elements on segments

You are given list of requests you need to execute. Eeach request is a triple (l,r,x), which means that to all the elements with indixes in the range [l,r) you need to add x. You are to print all the values in the array Input: n - size of the array (<10^5), m - the number of requests.

Example: Input: [0 2] [1 4] [5 6] Output: 1 2 2 1 1 1 1

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

int n, m;

// standard sweep line approach
void solution1() {
    cin >> n;
    cin >> m;
    vector< pair <int, int> > reqs;
    for (int i = 0 ; i < m; ++i) {
        int li,ri,x;
        cin >> li >> ri >> x;
        reqs.emplace_back(li, x);
        reqs.emplace_back(ri, -x);
    }
    sort(reqs.begin(), reqs.end());
    int j = 0;
    int sum = 0;
    for (int i = 0; i < n; ++i) {
        while (j < reqs.size() && i >= reqs[j].first) {
            sum += reqs[j].second;
            ++j;
        }
        cout << sum << " ";
    }
    cout << endl;
}
// simpler solution with prefix sums
void solution2() {
    cin >> n;
    cin >> m;
    vector< int > res(n);
    for (int i = 0 ; i < m; ++i) {
        int li,ri,x;
        cin >> li >> ri >> x;
        res[li] += x;
        res[ri] += -x;
    }
    int sum = 0; // to have for from 0
    for (int i = 0 ; i < n; ++i) {
        sum += res[i];
        res[i] = sum;
        cout << res[i] << " ";
    }
    cout << endl;
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(NULL);
    solution2(); 
    return 0;
}

Problem 2: paint fence

codeforces

You are given q intervals [l, r], you need to find subset of size q-2 such that the cumulative length of obtained segments is maximized. You know that l,r < 5e3.

So build prefix array, s.t. fence[i] -- how many segments are covering point i.

Try exclude worker by worker (so consider q-1 workers) and find the worker which gives the least contribution. For that, build another prefix array prefixSumD[i] -- how many points on [0,i] which are covered exactly by one segment. Than it is easy to see that to find number of sections segment x covers is prefixSumD[r[x]] - prefixSumD[l[x] - 1].

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

int fence[5001];

vector<pair<int,int>> pairs;

int main() {
    ios_base::sync_with_stdio(0); cin.tie(NULL);
    int n, q;
    cin >> n >> q;

    pairs.reserve(q);
    for (int i = 0; i < q; ++i) {
        int l,r;
        cin >> l >> r;
        pairs.emplace_back(l, r);
        while (l <= r) {
            ++fence[l];
            ++l;
        }
    }
    
    ++n; // because 1..n range
    int bestRes = 0;
    vector<int> prefix(n, 0); // array of prefix sums
    for (int i = 0; i < pairs.size(); ++i) {
        const auto[l,r] = pairs[i];
        int total = 0;
        int j = l;
        // delete this range
        while (j <= r) {
            --fence[j];
            ++j;
        }

        for (int v : fence)
            if (v != 0)
                ++total;

        // fill in prefix
        prefix[0] = fence[0] == 1? 1 : 0;
        for (j = 1; j < n; ++j) {
            prefix[j] = prefix[j-1] + (fence[j] == 1 ? 1 : 0);
        }

        int resMin = 1e6;
        for (j = 0; j < q; ++j) {
            if (j == i) continue;
            const auto[ll,rr] = pairs[j];
            assert(prefix[ll] <= prefix[rr]);
            resMin = min(resMin, prefix[rr] - (ll == 0? 0 :prefix[ll-1]));
        }
        int curRes = total - resMin;
        if (curRes > bestRes)
            bestRes = curRes;
        // add back the range
        j = l;
        while (j <= r) {
            ++fence[j];
            ++j;
        }
    }
    cout << bestRes << endl; 

    return 0;
}

Problem 3: Generalization on trees

cf.

Given a tree, and a list of pairs v_i, d_i, x_i add x_i to all the element in the subtree rooted at v_i of depth d_i.

// the main idea is that we simplify the problem to adding xi for the whole subtree rooted at vi
// to compensate extra increments, we subtract (-xi) to the leafs of the subree at vi with height di+1

#include <bits/stdc++.h>
using namespace std;
using lint = long long int;
vector<vector<int>> adj;
vector< vector<pair<lint,lint>> > w; // vi -> <di,xi>
vector<lint> res;

vector<lint> add; // this array allows to push forwards the updates for vertices at particular depth
// it stores increments for the vtx and when we achieve a vtx in dfs - it already
// contains information about increments from ancestors
// The way it work is similar to increment on subarray when in order to increment A[s..e] we do the following:
// A[s] += x, A[e] -= x and later do prefix sum. Here the prefix sum is argument of dfs (sum)

void dfs(int v, int p, int depth, lint sum)
{
    for (auto&[d, x] : w[v]) {
        add[depth] += x;
        if (depth + d + 1 < add.size()) {
            add[depth + d + 1] -= x;
        }
    }
    sum += add[depth];
    res[v] = sum;
    for (auto u : adj[v])
        if (u != p) {
            dfs(u, v, depth + 1, sum);
        }
    //restore state of add
    for (auto&[d, x] : w[v]) {
        add[depth] -= x;
        if (depth + d + 1 < add.size()) {
            add[depth + d + 1] += x;
        }
    }
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(NULL);

    int n;
    cin >> n;
    adj.resize(n);
    w.resize(n);
    res.resize(n);
    for (int i = 0; i < n-1; ++i) {
        int u,v;
        cin >> u >> v;
        --u; --v;
        adj[u].emplace_back(v);
        adj[v].emplace_back(u);
    }

    int m;
    cin >> m;
    for (int i = 0; i < m; ++i) {
        lint v, d, x;
        cin >> v >> d >> x;
        --v;
        w[v].emplace_back(d,x);
    }
    
    add.resize(3e5 + 9, 0);
    dfs(0, -1, 0, 0);

    for (lint v : res)
        cout << v << " ";
    cout << endl;

    return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment