Submission #1097593

#	Time	Username	Problem	Language	Result	Execution time	Memory
1097593		nguyen31hoang08minh2003	Addk (eJOI21_addk)	C++17	100 / 100	288 ms	12116 KiB

Main

#include <set>
#include <map>
#include <cmath>
#include <queue>
#include <deque>
#include <stack>
#include <math.h>
#include <bitset>
#include <vector>
#include <string>
#include <cstdio>
#include <cctype>
#include <numeric>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <iomanip>
#include <cassert>
#include <cstring>
#include <stdio.h>
#include <string.h>
#include <iterator>
#include <iostream>
#include <algorithm>
#include <strings.h>
#include <functional>
#include <unordered_set>
#include <unordered_map>
#define fore(i, a, b) for (int i = (a), i##_last = (b); i < i##_last; ++i)
#define fort(i, a, b) for (int i = (a), i##_last = (b); i <= i##_last; ++i)
#define ford(i, a, b) for (int i = (a), i##_last = (b); i >= i##_last; --i)
#define fi first
#define se second
#define pb push_back
#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
using namespace std;
using ll = long long;
using ld = long double;

template<class A, class B> bool maxi(A &a, const B &b) {return (a < b) ? (a = b, true):false;};
template<class A, class B> bool mini(A &a, const B &b) {return (a > b) ? (a = b, true):false;};

typedef unsigned long long ull;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<vi> vvi;
typedef vector<vii> vvii;

class SegmentTree {
private:

    int n;
    vector<ll> sum;

    ll query(int ql, int qr, int i, int l, int r) {
        if (r < ql || qr < l)
            return 0;
        if (ql <= l && r <= qr)
            return sum[i];
        const int m = l + r >> 1;
        return query(ql, qr, i << 1, l, m) + query(ql, qr, i << 1 | 1, m + 1, r);
    }

    void modify(int q, ll value, int i, int l, int r) {
        if (q < l || r < q)
            return;
        if (l == r) {
            sum[i] = value;
            return;
        }
        const int m = l + r >> 1;
        modify(q, value, i << 1, l, m);
        modify(q, value, i << 1 | 1, m + 1, r);
        sum[i] = sum[i << 1] + sum[i << 1 | 1];
    }

public:

    SegmentTree(): n(0) {};

    SegmentTree(const int n): n(n), sum(n + 5 << 2) {};

    void resize(const int m) {
        n = m;
        sum.resize(n + 5 << 2);
    }

    ll query(int ql, int qr) {
        if (ql > qr)
            return 0;
        return query(ql, qr, 1, 1, n);
    }

    void modify(int q, ll value) {
        modify(q, value, 1, 1, n);
    }

};

constexpr int MAX_N = 1E6 + 6;

int N, K, Q, A[MAX_N];
SegmentTree sum, increasedSum, decreasedSum;

ll findIncreasedSum(int l, int r) {
    /*

        A[l] * 1 + ... + A[r] * (r - l + 1)

    */
    return increasedSum.query(l, r) - sum.query(l, r) * (l - 1);
}

ll findDecreasedSum(int l, int r) {
    /*

        A[l] * (r - l + 1) + ... + A[r] * 1

    */
    return decreasedSum.query(l, r) - sum.query(l, r) * (N - r);
}

ll query(int l, int r, int m) {
    /*

        Example
            m = 5

            l = 0       l = 0       l = 0
            r = 4       r = 5       r = 9


            11111       122221      1234554321
            -----       -----|     0-----|||||
                         -----     1|-----||||
                                   2||-----|||
                                   3|||-----||
                                   4||||-----|
                                   5|||||-----
                                    0123456789


    */
    int L = l + m - 1, R = r - m + 1;
    if (L < R)
        return findIncreasedSum(l, L) + m * sum.query(L + 1, R - 1) + findDecreasedSum(R, r);

//    cerr << findIncreasedSum(l, R - 1) << ' ' << (R - l + 1) * sum.query(R, L) << ' ' << findDecreasedSum(L + 1, r) << '\n';

    return findIncreasedSum(l, R - 1) + (R - l + 1) * sum.query(R, L) + findDecreasedSum(L + 1, r);
}

void setValue(const int i, const int value) {
    A[i] = value;
    sum.modify(i, A[i]);
    increasedSum.modify(i, 1LL * i * A[i]);
    decreasedSum.modify(i, (N - i + 1LL) * A[i]);
}

void solveSecondQuery() {
    static int l, r, m;

    cin >> l >> r >> m;

    cout << query(l, r, m) << '\n';
}

void solveFirstQuery() {
    static int i[MAX_N], a[MAX_N];

    fort(k, 1, K)
        cin >> i[k];

    fore(k, 1, K)
        a[i[k]] = A[i[k + 1]];

    a[i[K]] = A[i[1]];

    fort(k, 1, K)
        setValue(i[k], a[i[k]]);
}

signed main() {

    int t;

    #ifdef LOCAL
    freopen("input.INP", "r", stdin);
//    freopen("output.OUT", "w", stdout);
    #endif // LOCAL
    cin.tie(0) -> sync_with_stdio(0);
    cout.tie(0);

    cin >> N >> K;

    sum.resize(N);
    increasedSum.resize(N);
    decreasedSum.resize(N);

    fort(i, 1, N) {
        cin >> A[i];
        setValue(i, A[i]);
    }

    cin >> Q;

    fore(_, 0, Q) {
        cin >> t;
        if (t == 2)
            solveSecondQuery();
        else if (t == 1)
            solveFirstQuery();
    }

    return 0;
}

/*

    Use three segment trees
        (one segment tree might not have to be used but its absence might make code to become more complicated)

*/

Compilation message (stderr)

Main.cpp: In member function 'll SegmentTree::query(int, int, int, int, int)':
Main.cpp:63:25: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   63 |         const int m = l + r >> 1;
      |                       ~~^~~
Main.cpp: In member function 'void SegmentTree::modify(int, ll, int, int, int)':
Main.cpp:74:25: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   74 |         const int m = l + r >> 1;
      |                       ~~^~~
Main.cpp: In constructor 'SegmentTree::SegmentTree(int)':
Main.cpp:84:43: warning: suggest parentheses around '+' inside '<<' [-Wparentheses]
   84 |     SegmentTree(const int n): n(n), sum(n + 5 << 2) {};
      |                                         ~~^~~
Main.cpp: In member function 'void SegmentTree::resize(int)':
Main.cpp:88:22: warning: suggest parentheses around '+' inside '<<' [-Wparentheses]
   88 |         sum.resize(n + 5 << 2);
      |                    ~~^~~

• Subtask 1: 36 / 36
• Subtask 2: 56 / 56
• Subtask 3: 8 / 8