Submission #675802

#TimeUsernameProblemLanguageResultExecution timeMemory
675802rafatoaIzbori (COCI22_izbori)C++17
110 / 110
579 ms16504 KiB
#include <bits/stdc++.h>
using namespace std;
#pragma GCC optimize ("Ofast")
 
#define F first
#define S second
#define vi vector<int>
#define vvi vector<vi>
#define pi pair<int, int>
#define vpi vector<pi>
#define vb vector<bool>
#define vvb vector<vb>
#define pb push_back
#define ppb pop_back
#define read(a) for(auto &x:a) cin >> x;
#define print(a) for(auto x:a) cout << x << " "; cout << "\n";
#define vc vector<char>
#define vvc vector<vc>
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
 
#define int long long
#define ld long double
const int INF = 1e18;
const int inf = 1e9;

struct FT{
    int n;
    vi tree;
    FT(int sz){
        n = sz;
        tree = vi(n+1);
    }

    void update(int k, int val){
        k++;
        while(k <= n){
            tree[k] += val;
            k += k&-k;
        }
    }

    int query(int l, int r){
        r++;
        int ans = 0;
        while(r >= 1){
            ans += tree[r];
            r -= r&-r;
        }
        while(l >= 1){
            ans -= tree[l];
            l -= l&-l;
        }
        return ans;
    }
};

void solve(){
    int n; cin >> n;
    vi a(n); read(a);

    map<int, vi> mp;
    for(int i=0; i<n; i++)
        mp[a[i]].pb(i);

    //Número de subarrays de longitud <= len en [l, r]
    auto f = [&](int l, int r, int len){
        if(l > r) return 0LL;
        int s = (r-l+1)*(r-l+2)/2;
        if(len < r-l+1) s -= (r-l+1-len)*(r-l+2-len)/2;
        return s;
    };

    int ans = 0;
    for(auto &[el, pos]:mp){
        if(pos.size() <= 1000){
            for(int i=0; i<pos.size(); i++)
            for(int j=i; j<pos.size(); j++){
                int mxlen = 2*(j-i)+1;
                if(pos[j]-pos[i]+1 > mxlen) continue;

                int l = (i ? pos[i-1]+1 : 0);
                int r = (j+1 < pos.size() ? pos[j+1]-1 : n-1);

                ans += f(l, r, mxlen)-f(pos[i]+1, r, mxlen)-f(l, pos[j]-1, mxlen)+
                    f(pos[i]+1, pos[j]-1, mxlen);
            }
        } else {
            FT ft(2*n+1);
            int curr = 0;
            ft.update(n, 1);
            for(int i=0; i<n; i++){
                curr += (a[i] == el ? 1 : -1);
                if(curr+n-1 >= 0) ans += ft.query(0, curr+n-1);
                ft.update(curr+n, 1);
            }
        }
    }
    cout << ans << "\n";
}
 
signed main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
 
    // #ifndef ONLINE_JUDGE
    //     freopen("in.txt", "r", stdin);
    //     freopen("out.txt", "w", stdout);
    // #endif

    int tt = 1;
    // cin >> tt;
    while(tt--)
        solve();
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void solve()':
Main.cpp:77:27: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   77 |             for(int i=0; i<pos.size(); i++)
      |                          ~^~~~~~~~~~~
Main.cpp:78:27: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   78 |             for(int j=i; j<pos.size(); j++){
      |                          ~^~~~~~~~~~~
Main.cpp:83:30: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   83 |                 int r = (j+1 < pos.size() ? pos[j+1]-1 : n-1);
      |                          ~~~~^~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...