Submission #1267035

#TimeUsernameProblemLanguageResultExecution timeMemory
1267035dzhoz0Izbori (COCI22_izbori)C++20
40 / 110
3094 ms14408 KiB
#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 = 400;
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;
    }
};


int bit[MAXN + 6];

void upd(int id, int x) {
    while(id <= MAXN + 1) {
        bit[id] += x;
        id += id & (-id);
    }
}

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

ll cntPositiveSubarraySum(vi &a) {
    int n = (int)a.size();
    vi ps(n); ps[0] = a[0];
    for(int i = 1; i < n; i++) ps[i] = ps[i - 1] + a[i];
    Compressor cmp;
    cmp.add(0);
    for(int i = 0; i < n; i++) cmp.add(ps[i]);
    cmp.init(); 
    
    upd(cmp.id(0), 1);

    ll ans = 0;
    for(int i = 0; i < n; i++) {
        ans += query(cmp.id(ps[i]) - 1);    
        upd(cmp.id(ps[i]), 1);
    }

    upd(cmp.id(0), -1);

    for(int i = 0; i < n; i++) {
        upd(cmp.id(ps[i]), -1);
    } 
    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);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...