Source Code

#include <bits/stdc++.h>
#define pb push_back
#define ll long long
#define pi pair<ll,ll>
using namespace std;
const ll inf = 1e5+9;
ll n,k;
map<ll,ll> cnt;
vector<pi> v;

void TestCases(){
    scanf("%lld%lld",&n,&k);
    ll ans = 0,isroot = 0;
    v.clear();
    cnt.clear();
    for(ll i=1;i<=n;i++){
        ll tmp;
        scanf("%lld",&tmp);
        if(tmp == 0)
            isroot++;
        else
            cnt[tmp]++;
    }
    if(isroot > 1)
        return void(puts("-1"));
    for(auto o:cnt)
        v.pb(o),ans += o.second;
    if(v.empty())
        return void(puts("1"));
    if(k == 1){
        if(v.size() == 1 && v[0] == pi(1,1))
            puts("2");
        else
            puts("-1");
        return ;
    }
    if(k == 2){
        vector<ll> tmp;
        ll last = 0;
        ans = 1;
        for(auto o:v){
            if(o.second > 2)
                return void(puts("-1"));
            ans += o.first - last,last = o.first;
            if(o.second == 2)
                tmp.pb(o.first);
        }
        last = 0;
        for(auto o:tmp)
            ans += o - last,last = o;
        printf("%lld\n",ans);
        return ;
    }
    if(v[0].first != 1)
        v.insert(v.begin(),pi(1,1)),ans++;
    for(ll i=v.size()-1;i>0;i--){
        while(v[i].first > v[i-1].first && v[i].second > 1){
            v[i].second= (v[i].second/(k-1) + (v[i].second%(k-1) > 0)),v[i].first--,
            ans += v[i].second;
        }

        if(v[i].first == v[i-1].first){
                if(v[i].second > v[i-1].second)
                    ans -= v[i-1].second,v[i-1].second = v[i].second;
            else
                ans -= v[i].second;
        }
        else
            ans += v[i].first - v[i-1].first - 1;
    }
    if(v[0].second > k)
        return void(puts("-1"));
    printf("%lld\n",ans+1);
}

int main(){
    ll t;
    scanf("%lld",&t);
    while(t--)
        TestCases();
}
Copy
Distances MNM
GNU G++17
297 ms
8.4 MB
Accepted