Submission #144414

#	Time	Username	Problem	Language	Result	Execution time	Memory
144414		xiaowuc1	Split the Attractions (IOI19_split)	C++14	100 / 100	161 ms	17016 KiB

split

#include <bits/stdc++.h>

using namespace std;

typedef pair<int, int> pii;

vector<int> edges[100005];

int ufpar[100005];
int ufsz[100005];

int find(int x) {
  return ufpar[x] == x ? x : (ufpar[x] = find(ufpar[x]));
}
bool merge(int x, int y) {
  x = find(x);
  y = find(y);
  if(x == y) return false;
  if(ufsz[x] < ufsz[y]) swap(x, y);
  ufsz[x] += ufsz[y];
  ufpar[y] = x;
  return true;
}

// spanning tree is rooted at 0
int treeSize[100005];
int dfsParent[100005];
bool dfsSeen[100005];
int assignedColor[100005];
int ret[100005];

bool ufcolor(int curr, int need, int want) {
  ret[curr] = want;
  if(--need == 0) return true;
  queue<int> q;
  q.push(curr);
  while(!q.empty()) {
    curr = q.front(); q.pop();
    for(int out: edges[curr]) {
      if(find(out) != find(curr)) continue;
      if(ret[out]) continue;
      ret[out] = want;
      if(--need == 0) return true;
      q.push(out);
    }
  }
  return false;
}

vector<int> consolidate(int n) {
  vector<int> ans;
  map<int, int> dp;
  for(int i = 0; i < n; i++) {
    assert(ret[i] >= 1 && ret[i] <= 3);
    ans.push_back(ret[i]);
    dp[ret[i]]++;
  }
  return ans;
}

void downColor(int curr, int par, int& need, int want) {
  assert(ret[curr] == 0);
  ret[curr] = want;
  assert(need-- > 0);
  if(need > 0 && dfsParent[curr] >= 0 && dfsParent[curr] != par) {
    downColor(dfsParent[curr], curr, need, want);
  }
  for(int out: edges[curr]) {
    if(need == 0) break;
    if(out == par) continue;
    if(dfsParent[out] != curr) continue;
    downColor(out, curr, need, want);
  }
}

bool upColor(int curr, int need, int want) {
  queue<int> q;
  q.push(curr);
  ret[curr] = want;
  if(--need == 0) return true;
  while(!q.empty()) {
    curr = q.front(); q.pop();
    for(int out: edges[curr]) {
      if(ret[out]) continue;
      ret[out] = want;
      if(--need == 0) return true;
      q.push(out);
    }
  }
  return false;
}

void fillRest(int n, int want) {
  for(int i = 0; i < n; i++) {
    if(ret[i] == 0) ret[i] = want;
  }
}

void dfs(int curr, int par) {
  treeSize[curr] = 1;
  dfsSeen[curr] = true;
  dfsParent[curr] = par;
  for(int out: edges[curr]) {
    if(dfsSeen[out]) continue;
    dfs(out, curr);
    treeSize[curr] += treeSize[out];
  }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
  for(int i = 0; i < p.size(); i++) {
    edges[p[i]].push_back(q[i]);
    edges[q[i]].push_back(p[i]);
  }
  dfs(0, -1);
  vector<pii> sizes;
  sizes.emplace_back(a, 1);
  sizes.emplace_back(b, 2);
  sizes.emplace_back(c, 3);
  sort(sizes.begin(), sizes.end());
  // is there a subtree such that it can be colored with size a and the other half can be colored with size b
  for(int i = 1; i < n; i++) {
    if(treeSize[i] >= sizes[0].first && n-treeSize[i] >= sizes[1].first) {
      downColor(i, dfsParent[i], sizes[0].first, sizes[0].second);
      assert(upColor(dfsParent[i], sizes[1].first, sizes[1].second));
      fillRest(n, sizes[2].second);
      return consolidate(n);
    }
    if(treeSize[i] >= sizes[1].first && n-treeSize[i] >= sizes[0].first) {
      downColor(dfsParent[i], i, sizes[0].first, sizes[0].second);
      assert(upColor(i, sizes[1].first, sizes[1].second));
      fillRest(n, sizes[2].second);
      return consolidate(n);
    }
  }
  // find a node such that, when deleted, all trees have size < a
  int critical = -1;
  for(int i = 0; i < n; i++) {
    int largestTree = 0;
    int parentSize = n-1;
    for(int out: edges[i]) {
      if(dfsParent[out] != i) continue;
      largestTree = max(largestTree, treeSize[out]);
      parentSize -= treeSize[out];
    }
    largestTree = max(largestTree, parentSize);
    if(largestTree >= sizes[0].first) continue;
    critical = i;
    break;
  }
  assert(critical != -1);
  for(int i = 0; i < n; i++) {
    ufpar[i] = i;
    ufsz[i] = 1;
  }
  for(int i = 1; i < n; i++) {
    if(i == critical || dfsParent[i] == critical) continue;
    assert(merge(i, dfsParent[i]));
    assert(ufsz[find(i)] < sizes[0].first);
  }
  for(int i = 0; i < n; i++) {
    for(int out: edges[i]) {
      if(i == critical || out == critical) continue;
      if(!merge(i, out)) continue;
      if(ufsz[find(i)] >= sizes[0].first) {
        ufcolor(i, sizes[0].first, sizes[0].second);
        assert(upColor(critical, sizes[1].first, sizes[1].second));
        fillRest(n, sizes[2].second);
        return consolidate(n);
      }
    }
  }
  vector<int> ans;
  ans.resize(n);
  return ans;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:111:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i = 0; i < p.size(); i++) {
                  ~~^~~~~~~~~~

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