Submission #295554

#	Time	Username	Problem	Language	Result	Execution time	Memory
295554		tzero	Split the Attractions (IOI19_split)	C++14	29 / 100	112 ms	11768 KiB

split

#include <vector>
#include <iostream>
#include <algorithm>
#include <cassert>
using namespace std;
const int maxn = (int)1e5+7;
int color_idx;
vector<int> adj[maxn];
vector<int> ans, sub;
vector<pair<int,int> > colors;

void dfs1(int a) {
  ans[a] = colors[color_idx].second;
  colors[color_idx].first--; if(colors[color_idx].first == 0) color_idx++;
  for(int b : adj[a]) if(ans[b] == 0) dfs1(b);
}

void dfs2(int a, int p) {
  sub[a] = 1;
  for(int b : adj[a]) if(b != p) {
    dfs2(b, a); sub[a] += sub[b];
  }
}

void dfs3(int a) {
  ans[a] = colors[color_idx].second;
  colors[color_idx].first--;

  for(int b : adj[a]) if(colors[color_idx].first > 0 && ans[b] == 0 && sub[b] < sub[a]) {
    dfs3(b);
  }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
  for(int i = 0; i < (int)q.size(); i++) {
    adj[q[i]].push_back(p[i]);  
    adj[p[i]].push_back(q[i]);
  }

  ans.assign(n, 0);

  bool two_max = true;
  for(int i = 0; i < n; i++) {
    if((int)adj[i].size() > 2) {
      two_max = false; 
    }
  }

  colors.push_back(make_pair(a, 1));
  colors.push_back(make_pair(b, 2));
  colors.push_back(make_pair(c, 3));
  sort(colors.begin(), colors.end());

  if(two_max) {
    bool found = false; color_idx = 0;
    for(int i = 0; found && i < n; i++) {
      if((int)adj[i].size() == 1) { found = true; dfs1(i); }
    }
    if(!found) dfs1(0);
    assert(colors[0].first + colors[1].first + colors[2].first == 0); 
    return ans;

  } else if((int)q.size() == n - 1) { // tree
    bool found = false;
    sub.resize(n);
    dfs2(0, 0); // subtree size

    for(int i = 0; !found && i < n; i++) {
      if(sub[i] >= colors[0].first && n - sub[i] >= colors[1].first) {
        found = true;   
        color_idx = 0;
        dfs3(i);
        color_idx = 1;
        dfs3(0);
      } else if(sub[i] >= colors[1].first && n - sub[i] >= colors[0].first) {
        found = true;
        color_idx = 1;
        dfs3(i);
        color_idx = 0;
        dfs3(0);
      }
    }

    if(found) {
      for(int i = 0; i < n; i++) if(ans[i] == 0) { 
        ans[i] = colors[2].second;  
      }
    }

    return ans;
  } else {
    return ans;
  }
}

int my_main() { 
  int n, m;
  cin >> n >> m;
  int a, b, c;
  cin >> a >> b >> c;
  vector<int> p(m), q(m);
  for(int i = 0; i < m; i++) {
    cin >> p[i] >> q[i];
  }
  vector<int> res = find_split(n, a, b, c, p, q);
  for(int x : res) cout << x << " ";
  cout << endl;
}

Compilation message (stderr)

split.cpp: In function 'int my_main()':
split.cpp:108:1: warning: no return statement in function returning non-void [-Wreturn-type]
  108 | }
      | ^

• Subtask 1: 7 / 7
• Subtask 2: 0 / 11
• Subtask 3: 22 / 22
• Subtask 4: 0 / 24
• Subtask 5: 0 / 36