Submission #2547212


Source Code Expand

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

struct UnionFind{
  Int n;
  vector<Int> r,p;
  UnionFind(){}
  UnionFind(Int sz):n(sz),r(sz,1),p(sz,0){iota(p.begin(),p.end(),0);}
  Int find(Int x){
    return (x==p[x]?x:p[x]=find(p[x]));
  }
  bool same(Int x,Int y){
    return find(x)==find(y);
  }
  void unite(Int x,Int y){
    x=find(x);y=find(y);
    if(x==y) return;
    if(r[x]<r[y]) swap(x,y);
    r[x]+=r[y];
    p[y]=x;
  }
  Int size(Int x){
    return r[find(x)];
  }
};

//INSERT ABOVE HERE
signed main(){
  Int n,m,k;
  cin>>n>>m>>k;
  vector<Int> c(n);
  for(Int i=0;i<n;i++) cin>>c[i];
  using T = tuple<Int, Int, Int>;
  vector<T> v;
  for(Int i=0;i<m;i++){
    Int a,b,w;
    cin>>a>>b>>w;
    a--;b--;
    v.emplace_back(w,a,b);
  }
  sort(v.begin(),v.end());

  UnionFind uf(n+k);
  for(Int i=0;i<n;i++){
    c[i]--;
    if(~c[i]) uf.unite(i,n+c[i]);
  }
  
  queue<Int> q;
  for(Int i=0;i<k;i++)
    if(uf.size(n+i)==1)
      q.emplace(i);  
  
  Int ans=0,cnt=1;
  priority_queue<Int> pq;
  for(Int i=0;i<m;i++){
    if(uf.size(0)==n+k) break;
    Int w,a,b;
    tie(w,a,b)=v[i];
    if(uf.same(a,b)) continue;
    //cout<<w<<" "<<a<<" "<<b<<endl;
    ans+=w;
    pq.emplace(w);
    uf.unite(a,b);
    if(c[a]<0){
      if(!q.empty()){
	c[a]=q.front();q.pop();
	uf.unite(a,n+c[a]);
      }else cnt++;
    }
    if(c[b]<0){
      if(!q.empty()){
	c[b]=q.front();q.pop();
	uf.unite(b,n+c[b]);
      }else cnt++;
    }
    //cout<<cnt<<" "<<uf.size(0)<<endl;
  }
  //cout<<cnt<<" "<<uf.size(0)<<endl;
  for(Int i=0;i<n+k;i++)
    if(uf.find(i)==i) cnt--;
  
  while(cnt>1){
    //cout<<pq.top()<<endl;
    ans-=pq.top();pq.pop();
    cnt--;
  }
  //cout<<ans<<" "<<cnt<<endl;
  if(cnt<0) ans=-1;
  cout<<ans<<endl;
  return 0;
}

Submission Info

Submission Time
Task A - Colorful MST
User beet
Language C++14 (GCC 5.4.1)
Score 200
Code Size 1864 Byte
Status WA
Exec Time 134 ms
Memory 9072 KB

Judge Result

Set Name Sample Subtask1 Subtask2 Subtask3 Subtask4
Score / Max Score 0 / 0 100 / 100 100 / 100 0 / 200 0 / 300
Status
AC × 4
AC × 7
AC × 12
AC × 2
WA × 5
AC × 18
WA × 10
Set Name Test Cases
Sample 00_example_01.txt, 00_example_02.txt, 00_example_03.txt, 00_example_04.txt
Subtask1 00_example_03.txt, s1_01.txt, s1_02.txt, s1_03.txt, s1_04.txt, s1_05.txt, s1_06.txt
Subtask2 s1_01.txt, s1_02.txt, s1_03.txt, s1_04.txt, s1_05.txt, s1_06.txt, s2_07.txt, s2_08.txt, s2_09.txt, s2_10.txt, s2_11.txt, s2_12.txt
Subtask3 00_example_02.txt, s3_13.txt, s3_14.txt, s3_15.txt, s3_16.txt, s3_17.txt, s3_18.txt
Subtask4 00_example_01.txt, 00_example_02.txt, 00_example_03.txt, 00_example_04.txt, s1_01.txt, s1_02.txt, s1_03.txt, s1_04.txt, s1_05.txt, s1_06.txt, s2_07.txt, s2_08.txt, s2_09.txt, s2_10.txt, s2_11.txt, s2_12.txt, s3_13.txt, s3_14.txt, s3_15.txt, s3_16.txt, s3_17.txt, s3_18.txt, s4_19.txt, s4_20.txt, s4_21.txt, s4_22.txt, s4_23.txt, s4_24.txt
Case Name Status Exec Time Memory
00_example_01.txt AC 1 ms 256 KB
00_example_02.txt AC 1 ms 256 KB
00_example_03.txt AC 1 ms 256 KB
00_example_04.txt AC 1 ms 256 KB
s1_01.txt AC 1 ms 256 KB
s1_02.txt AC 131 ms 8688 KB
s1_03.txt AC 1 ms 256 KB
s1_04.txt AC 1 ms 256 KB
s1_05.txt AC 73 ms 5876 KB
s1_06.txt AC 130 ms 8176 KB
s2_07.txt AC 18 ms 1276 KB
s2_08.txt AC 134 ms 7792 KB
s2_09.txt AC 94 ms 3828 KB
s2_10.txt AC 88 ms 3828 KB
s2_11.txt AC 122 ms 5232 KB
s2_12.txt AC 121 ms 5104 KB
s3_13.txt WA 16 ms 1404 KB
s3_14.txt AC 120 ms 8944 KB
s3_15.txt WA 97 ms 5232 KB
s3_16.txt WA 90 ms 6000 KB
s3_17.txt WA 121 ms 9072 KB
s3_18.txt WA 124 ms 7536 KB
s4_19.txt WA 18 ms 1404 KB
s4_20.txt WA 126 ms 6640 KB
s4_21.txt WA 98 ms 4208 KB
s4_22.txt WA 93 ms 5232 KB
s4_23.txt AC 130 ms 8560 KB
s4_24.txt WA 127 ms 7280 KB