Submission #1326694

#TimeUsernameProblemLanguageResultExecution timeMemory
1326694SmuggingSpunSplit the Attractions (IOI19_split)C++20
100 / 100
97 ms26156 KiB
#include "split.h"
#include<bits/stdc++.h>
using namespace std;
const int lim = 1e5 + 5;
int n, m, a, b, c;
vector<int>u, v, g[lim];
namespace sub1{
  bool work_subtask(){
    for(int i = 0; i < n; i++){
      if(g[i].size() > 2){
        return false;
      }
    }
    return true;
  }
  vector<int>solve(){
    vector<int>p;
    vector<bool>vis(n, false);
    function<void(int)>dfs;
    dfs = [&] (int s){
      p.push_back(s);
      vis[s] = true;
      for(int& i : g[s]){
        int d = u[i] ^ v[i] ^ s;
        if(!vis[d]){
          dfs(d);
        }
      }
    };
    for(int i = 0; i < n; i++){
      if(g[i].size() == 1){
        dfs(i);
        break;
      }
    }
    if(p.empty()){
      dfs(0);
    }
    vector<int>ans(n, 3);
    for(int i = 0; i < a; i++){
      ans[p[i]] = 1;
    }
    for(int i = 0; i < b; i++){
      ans[p[a + i]] = 2;
    }
    return ans;
  }
}
namespace sub2{
  vector<int>solve(){
    vector<int>ans(n);
    vector<bool>vis(n, false);
    function<void(int)>dfs;
    dfs = [&] (int s){
      vis[s] = true;
      if(b > 0){
        ans[s] = 2;
        b--;
      }
      else if(c > 0){
        ans[s] = 3;
        c--;
      }
      else{
        ans[s] = 1;
      }
      for(int& i : g[s]){
        int d = u[i] ^ v[i] ^ s;
        if(!vis[d]){
          dfs(d);
        }
      }
    };
    dfs(0);
    return ans;
  }
}
int IDX[3] = {1, 2, 3};
namespace sub3{
  vector<int>solve(){
    vector<int>siz(n), par(n), ans(n, 0);
    function<void(int)>dfs_1;
    dfs_1 = [&] (int s){
      siz[s] = 1;
      for(int& i : g[s]){
        int d = u[i] ^ v[i] ^ s;
        if(d != par[s]){
          par[d] = s; 
          dfs_1(d);
          siz[s] += siz[d];
        }
      }
    };
    dfs_1(par[0] = 0);
    function<void(int, int, int&, const int)>dfs_2;
    dfs_2 = [&] (int s, int p, int& cnt, const int color){
      if(cnt > 0){
        cnt--;
        ans[s] = color;
      }
      for(int& i : g[s]){
        int d = u[i] ^ v[i] ^ s;
        if(d != p){
          dfs_2(d, s, cnt, color);
        }
      }
    };
    for(int i = 1; i < n; i++){
      if(siz[i] >= a && n - siz[i] >= b){
        fill(ans.begin(), ans.end(), IDX[2]);
        dfs_2(i, par[i], a, IDX[0]);
        dfs_2(par[i], i, b, IDX[1]);
        break;
      }
      else if(siz[i] >= b && n - siz[i] >= a){
        fill(ans.begin(), ans.end(), IDX[2]);
        dfs_2(i, par[i], b, IDX[1]);
        dfs_2(par[i], i, a, IDX[0]);
        break;
      }
    }
    return ans;
  }
}
namespace sub45{
  vector<int>back_edge, ans, tree[lim];
  vector<vector<int>>part;
  int root = 0, par[lim], siz[lim], color[lim];
  void dfs_1(int s){
    siz[s] = 1;
    for(int& i : g[s]){
      int d = u[i] ^ v[i] ^ s;
      if(d != par[s] && par[d] == -1){
        tree[par[d] = s].push_back(i);
        tree[d].push_back(i);
        dfs_1(d);
        siz[s] += siz[d];
      }
      else if(d != par[s]){
        back_edge.push_back(i);
      }
    }
  }
  void dfs_2(int s, int p){
    part[color[s]].push_back(s);
    for(int& i : tree[s]){
      int d = u[i] ^ v[i] ^ s;
      if(d != p){
        color[d] = color[s]; 
        dfs_2(d, s);
      }
    }
  }
  void dfs_3(int s, const int x){
    if(a > 0){
      a--;
      ans[s] = IDX[0];
    }
    color[s] = -1;
    for(int& i : g[s]){
      int d = u[i] ^ v[i] ^ s;
      if(color[d] == x){
        dfs_3(d, x);
      }
    }
  }
  void dfs_4(int s){
    if(b > 0){
      b--;
      ans[s] = IDX[1];
    }
    int x = color[s];
    color[s] = -1;
    for(int& i : g[s]){
      int d = u[i] ^ v[i] ^ s;
      if(color[d] != -1){
        dfs_4(d);
      }
    }
  }
  bool work(int id){
    if(part[id].size() < a){
      return false;
    }
    ans.assign(n, IDX[2]);
    dfs_3(part[id][0], id);
    dfs_4(root);
    return true;
  }
  vector<int>solve(){
    memset(par, -1, sizeof(par));
    dfs_1(par[0] = 0);
    while(true){
      bool flag = true;
      for(int& i : tree[root]){
        int d = u[i] ^ v[i] ^ root;
        if(d != par[root] && siz[d] > (n >> 1)){
          root = d;
          flag = false;
          break;
        }
      }
      if(flag){
        break;
      }
    }
    for(int& i : tree[root]){
      int d = u[i] ^ v[i] ^ root;
      color[d] = part.size();
      part.push_back(vector<int>());
      dfs_2(d, root);
    }
    color[root] = -2;
    for(int i = 0; i < part.size(); i++){
      if(work(i)){
        return ans;
      }
    }
    for(int& i : back_edge){
      if(color[u[i]] != color[v[i]] && u[i] != root && v[i] != root){
        int large_color = part[color[u[i]]].size() > part[color[v[i]]].size() ? color[u[i]] : color[v[i]], small_color = color[u[i]] ^ color[v[i]] ^ large_color;
        for(int& j : part[small_color]){
          part[color[j] = large_color].push_back(j);
        }
        part[small_color].clear();
        if(work(large_color)){
          return ans;
        }
      }
    }
    return vector<int>(n, 0);
  }
}
vector<int>find_split(int _n, int _a, int _b, int _c, vector<int>_u, vector<int>_v){
  n = _n;
  a = _a;
  b = _b;
  c = _c;
  swap(u, _u);
  swap(v, _v);
  m = u.size();
  for(int i = 0; i < m; i++){
    g[u[i]].push_back(i);
    g[v[i]].push_back(i);
  }
  if(sub1::work_subtask()){
    return sub1::solve();
  }
  if(a == 1){
    return sub2::solve();
  }
  if(a > b){
    swap(a, b);
    swap(IDX[0], IDX[1]);
  }
  if(b > c){
    swap(b, c);
    swap(IDX[1], IDX[2]);
  }
  if(a > b){
    swap(a, b);
    swap(IDX[0], IDX[1]);
  }
  if(m == n - 1){
    return sub3::solve();
  }
  return sub45::solve();
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...