Submission #1267150

#	Time	Username	Problem	Language	Result	Execution time	Memory
1267150		dzhoz0	Izbori (COCI22_izbori)	C++20	110 / 110	1307 ms	14404 KiB

Main

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

#define ll long long
#define INF 1e18
#define f first
#define s second
#define pii pair<int, int>
#define vi vector<int>

const int MOD = 1'000'000'000 + 7;

void setIO(string name = "")
{
    ios_base::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
#ifdef LOCAL
    freopen("inp.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
#else
    if (!name.empty())
    {
        freopen((name + ".INP").c_str(), "r", stdin);
        freopen((name + ".OUT").c_str(), "w", stdout);
    }
#endif
}

ll countPairs(int l, int r, int u, int v, int k) {
    int R = r - 1; 
    int mix = l + 1;
    int mxx = u;
    if (mix > mxx || R < v) return 0;

    int x0 = max(mix, v - k + 1);
    if (x0 > mxx) return 0;

    ll ans = 0;

    int xt = R - k + 1;
    int x1 = min(mxx, xt);

    if (x0 <= x1) {
        ll n1 = x1 - x0 + 1;
        ll sum_x = (x0 + x1) * n1 / 2;
        ans += sum_x + n1 * (k - v);
    }

    int x2_st = max(x0, x1 + 1);
    if (x2_st <= mxx) {
        ll n2 = mxx - x2_st + 1;
        ll per = R - v + 1;
        if (per > 0) ans += n2 * per;
    }

    return ans;
}

const int MAXN = 2e5;
const int B = 500;
int n;
int a[MAXN + 5];
map<int, vi> mp;

struct Compressor {
    vi v;
    int sz() { return (int)v.size(); }

    void add(int x) {
        v.push_back(x);
    }
    
    void init() {
        sort(v.begin(), v.end());
        v.erase(unique(v.begin(), v.end()), v.end());
    }

    int id(int key) {
        return lower_bound(v.begin(), v.end(), key) - v.begin() + 1;
    }
};




ll cntPositiveSubarraySum(vi &a) {
    int n = (int)a.size();
    int sz = 2 * n + 5;
    vi bit(sz + 5, 0); 
    int offset = n + 2;

    auto upd = [&](int id) {
        while(id <= sz) {
            bit[id]++;
            id += id & (-id);
        }
    };

    auto query = [&](int id) -> int {
        int ans = 0;
        while(id > 0) {
            ans += bit[id];
            id -= id & (-id);
        }
        return ans;
    };


    ll ans = 0;
    int cur = 0;
    
    upd(cur + offset);

    for(int i = 0; i < n; i++) {
        cur += a[i];
        ans += query(cur + offset - 1);
        upd(cur + offset);
    }

    return ans;

}


void solve()
{
    cin >> n;
    for(int i = 1; i <= n; i++) {
        cin >> a[i];
        mp[a[i]].push_back(i);
    }
    ll ans = 0;
    for(auto &[val, p] : mp) {
        if((int)p.size() <= B) {
            for(int i = 0; i < (int)p.size(); i++) {
                for(int j = i; j < (int)p.size(); j++) {
                    int l = (i == 0 ? 0 : p[i - 1]);
                    int r = (j == (int)p.size() - 1 ? n + 1 : p[j + 1]);
                    int u = p[i];
                    int v = p[j];
                    int k = 2 * (j - i + 1) - 1;
                    ans += countPairs(l, r, u, v, k);
                }
            }
        }
        else {
            vi b(n, -1);
            for(int x : p) b[x - 1] = 1;
            ans += cntPositiveSubarraySum(b);
        }
    }
    
    cout << ans << '\n';
    
}
signed main()
{
    setIO();
    int t = 1;
    // cin >> t;
    while (t--)
        solve();
}

Compilation message (stderr)

Main.cpp: In function 'void setIO(std::string)':
Main.cpp:23:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   23 |         freopen((name + ".INP").c_str(), "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:24:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   24 |         freopen((name + ".OUT").c_str(), "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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