Submission #765022

#	Time	Username	Problem	Language	Result	Execution time	Memory
765022		The_Samurai	Izbori (COCI22_izbori)	C++17	40 / 110	3065 ms	5084 KiB

Main

#include "bits/stdc++.h"

using namespace std;
using ll = long long;

ll sum(ll l, ll r) {
    // l^2 + ... + r^2
    l--;
    return r * (r + 1) / 2 - l * (l + 1) / 2;
    return r * (r + 1) * (2 * r + 1) / 6 - l * (l + 1) * (2 * l + 1) / 6;
}

int brute(int n, vector<int> &a) {
    int ans = 0;
    for (int i = 0; i < n; i++) {
        vector<int> cnt(n);
        int mx = 0;
        for (int j = i; j < n; j++) {
            cnt[a[j]]++;
            mx = max(mx, cnt[a[j]]);
            if (2 * mx > (j - i + 1)) {
                ans++;
            }
        }
    }
    return ans;
}

void solve() {
    int n;
    cin >> n;
    map<int, int> mp;
    vector<int> cnt, a(n);
    {
        for (int &x: a) {
            cin >> x;
            mp[x];
        }
        int z = 0;
        for (auto &it: mp) {
            it.second = z++;
        }
        cnt.assign(mp.size(), 0);
        for (int &x: a) {
            x = mp[x];
            cnt[x]++;
        }
    }
    vector<vector<int>> pos(mp.size());
    for (int i = 0; i < n; i++) {
        pos[a[i]].emplace_back(i);
    }
    int N = sqrt(N);
    ll ans = 0;
    for (int i = 0; i < mp.size(); i++) {
        if (cnt[i] <= N) {
            for (int l = 0; l < cnt[i]; l++) {
                for (int r = l; r < cnt[i]; r++) {
                    int left = 0, right = n - 1, plus = r - l + 1;
                    int have = 2 * plus - (pos[i][r] - pos[i][l] + 1);
                    if (have <= 0) {
                        continue;
                    }
                    if (l > 0) {
                        left = pos[i][l - 1] + 1;
                    }
                    if (r + 1 < cnt[i]) {
                        right = pos[i][r + 1] - 1;
                    }
                    right = min(right, pos[i][r] + (have - 1));
                    left = max(left, pos[i][l] - (have - 1));
                    int right2 = left + (2 * plus - 1) - 1;
                    int left2 = right - (2 * plus - 1) + 1;
                    if (right2 < right) {
                        int x = pos[i][l] - left2 + 1;
                        ans += sum(x, right - right2 + x - 1);
                    } else {
                        right2 = right;
                    }
                    ans += (pos[i][l] - left + 1) * (right2 - pos[i][r] + 1);
                }
            }
        } else {
            vector<int> pref(n + 1), cs(2 * n + 1);
            for (int j = 0; j < n; j++) {
                pref[j + 1] = pref[j] + (a[j] == i ? 1 : -1);
            }
            for (int j = 0; j <= n; j++) {
                pref[j] += n;
            }
            cs[pref[n]]++;
            int now = 0;
            for (int j = n - 1; j >= 0; j--) {
                // count of k that pref[k] > pref[j] and k > j
                if (pref[j] < pref[j + 1]) {
                    now += cs[pref[j + 1]];
                } else {
                    now -= cs[pref[j]];
                }
                ans += now;
                cs[pref[j]]++;
            }
        }
    }
#ifdef test_cases
    cout << brute(n, a) << '\n';
#endif
    cout << ans;
}

int main() {
    ios_base::sync_with_stdio(false);
    cout.tie(nullptr);
    cin.tie(nullptr);

    int queries = 1;
#ifdef test_cases
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
    cin >> queries;
#else
    //    cin >> queries;
#endif

    for (int test_case = 1; test_case <= queries; test_case++) {
#ifdef test_cases
        cout << "Test case: " << test_case << '\n';
#endif
        solve();
        cout << '\n';
    }
}

Compilation message (stderr)

Main.cpp: In function 'void solve()':
Main.cpp:55:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::map<int, int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   55 |     for (int i = 0; i < mp.size(); i++) {
      |                     ~~^~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 15 / 15
• Subtask 3: 15 / 15
• Subtask 4: 0 / 70