Cod sursa(job #2355573)

Utilizator stefdascalescuStefan Dascalescu stefdascalescu Data 26 februarie 2019 10:08:55
Problema Ferma Scor 20
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 3.44 kb
#include<bits/stdc++.h>
using namespace std;
ifstream f("ferma.in");
ofstream g("ferma.out");
int n, k;
int ans = 0, s;
int v[10002], ord[10002], sum[10002];
int tt[10002];
int rg[10002];
bool was[10002];
multiset<int>sume, bad;
bool cmp(int a, int b)
{
    return v[a] > v[b];
}
int Find(int a)
{
    if(tt[a] == a)
        return a;
    return tt[a] = Find(tt[a]);
}
void Union(int a, int b)
{
    if(Find(a) == Find(b))
        return;
    a = Find(a), b = Find(b);
    multiset<int> ::iterator it;
    it = sume.lower_bound(sum[Find(a)]);
    if(*it == sum[Find(a)])
        sume.erase(it), s -= *it;
    else
    {
        it = bad.lower_bound(sum[Find(a)]);
        bad.erase(it);
    }
    it = sume.lower_bound(sum[Find(b)]);
    if(*it == sum[Find(b)])
        sume.erase(it), s -= *it;
    else
    {
        it = bad.lower_bound(sum[Find(b)]);
        bad.erase(it);
    }
    sume.insert(sum[Find(a)] + sum[Find(b)]);
    s += sum[Find(a)] + sum[Find(b)];
    if(rg[a] > rg[b])
    {
        tt[b] = a;
        rg[a] += rg[b];
        sum[a] += sum[b];
    }
    else
    {
        tt[a] = b;
        rg[b] += rg[a];
        sum[b] += sum[a];
    }
}
int main()
{
    f >> n >> k;
    for(int i = 1; i <= n; ++i)
        f >> v[i], ord[i] = i, tt[i] = i, rg[i] = 1, sum[i] = v[i];
    sort(ord+1, ord+n+1, cmp);
    for(int i = 1; i <= n; ++i)
    {
        int val = ord[i];
        was[val] = 1;
        sume.insert(v[val]);
        s += v[val];
        while(!bad.empty())
        {
            multiset<int> ::iterator it = bad.end();
            --it;
            if(sume.size() >= k && *it <= *sume.begin())
                break;
            sume.insert(*it);
            s += *it;
            bad.erase(it);
        }
        while(sume.size() > k)
        {
            multiset<int> ::iterator it = sume.begin();
            s -= *it;
            bad.insert(*it);
            sume.erase(it);
        }
        if(i >= k)
            ans = max(ans, s);
        if(val > 1 && was[val - 1])
            Union(Find(val - 1), Find(val));
        else
            if(val == 1 && was[n])
                Union(Find(n), Find(val));
        while(!bad.empty())
        {
            multiset<int> ::iterator it = bad.end();
            --it;
            if(sume.size() >= k && *it <= *sume.begin())
                break;
            sume.insert(*it);
            s += *it;
            bad.erase(it);
        }
        while(sume.size() > k)
        {
            multiset<int> ::iterator it = sume.begin();
            s -= *it;
            bad.insert(*it);
            sume.erase(it);
        }
        if(i >= k)
            ans = max(ans, s);
        if(val < n && was[val + 1])
            Union(Find(val), Find(val + 1));
        else
            if(val == n && was[1])
                Union(Find(val), Find(1));
        while(!bad.empty())
        {
            multiset<int> ::iterator it = bad.end();
            --it;
            if(sume.size() >= k && *it <= *sume.begin())
                break;
            sume.insert(*it);
            s += *it;
            bad.erase(it);
        }
        while(sume.size() > k)
        {
            multiset<int> ::iterator it = sume.begin();
            s -= *it;
            bad.insert(*it);
            sume.erase(it);
        }
        if(i >= k)
            ans = max(ans, s);
    }
    g << ans;
    return 0;
}