Problem p02749 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<ll, ll> ii;
typedef tuple<ll, ll, ll> iii;
typedef vector<ll> vi;
typedef vector<ii> vii;
typedef vector<iii> viii;
typedef vector<vi> vvi;
typedef vector<vii> vvii;
#define REP(i, n) for (ll i = 0; i < n; ++i)
#define REPR(i, n) for (ll i = n - 1; i >= 0; --i)
#define FOR(i, m, n) for (ll i = m; i < n; ++i)
#define FORR(i, m, n) for (ll i = n - 1; i >= m; --i)
#define FORE(x, xs) for (const auto &x : xs)
#define FORI(i, v) for (auto i = v.begin(); i != v.end(); i++)
#define ALL(v) v.begin(), v.end()
#define CHMIN(x, y) x = min(x, y)
#define CHMAX(x, y) x = max(x, y)
#define YES(b) cout << ((b) ? "YES" : "NO") << endl
#define Yes(b) cout << ((b) ? "Yes" : "No") << endl
#define DOUBLE(d) cout << fixed << setprecision(15) << (d) << endl

struct UnionFind {
private:
  int N;
  vi parent;
  vi rank_;

  void init(int n) {
    N = n;
    parent.resize(0);
    parent.clear();
    parent.resize(n);
    REP(i, N) parent[i] = i;
    rank_.resize(0);
    rank_.clear();
    rank_.resize(n);
  }

public:
  UnionFind() { init(0); }
  UnionFind(int n) { init(n); }

  int root(int x) { return (parent[x] == x) ? x : parent[x] = root(parent[x]); }

  bool isSame(int x, int y) { return root(x) == root(y); }

  void merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return;

    if (rank_[x] < rank_[y])
      parent[x] = y;
    else {
      parent[y] = x;
      if (rank_[x] == rank_[y])
        rank_[x]++;
    }
  }
};

class BipartiteGraph {
private:
  int N;
  UnionFind uf;
  bool valid;

  void init(int n) {
    N = n;
    uf = UnionFind(2 * N);
    valid = true;
  }

public:
  BipartiteGraph() { init(0); }
  BipartiteGraph(int n) { init(n); }

  void add(int u, int v) {
    uf.merge(u, v + N);
    uf.merge(u + N, v);

    valid &= !(uf.isSame(u, u + N) || uf.isSame(v, v + N));
  }

  bool isValid() { return valid; }

  pair<set<int>, set<int>> get() {
    assert(valid);

    set<int> u, v;
    REP(i, N) {
      if (uf.isSame(i, 0))
        u.insert(i);
      else
        v.insert(i);
    }
    return make_pair(u, v);
  }
};

int N;
BipartiteGraph bi;

int main() {
  cin >> N;
  bi = BipartiteGraph(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    bi.add(a, b);
  }

  vi ans(N);
  set<int> u, v;
  tie(u, v) = bi.get();

  if (v.size() > u.size())
    swap(u, v);
  int n = 1;
  FORI(iter, u) {
    ans[*iter] = n;
    n += 3;
    if (n > N)
      break;
  }
  n = 2;
  if (v.size() >= (N + 1) / 3) {
    FORI(iter, v) {
      ans[*iter] = n;
      n += 3;
      if (n > N)
        break;
    }
  } else {
    for (auto iter = next(u.begin(), (N + 2) / 3);; iter++) {
      FORI(iter, v) {
        ans[*iter] = n;
        n += 3;
        if (n > N)
          break;
      }
    }
  }

  n = 3;
  REP(i, N) {
    if (ans[i] != 0)
      continue;
    ans[i] = n;
    n += 3;
  }

  REP(i, N) cout << ans[i] << " ";
  cout << endl;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<ll, ll> ii;
typedef tuple<ll, ll, ll> iii;
typedef vector<ll> vi;
typedef vector<ii> vii;
typedef vector<iii> viii;
typedef vector<vi> vvi;
typedef vector<vii> vvii;
#define REP(i, n) for (ll i = 0; i < n; ++i)
#define REPR(i, n) for (ll i = n - 1; i >= 0; --i)
#define FOR(i, m, n) for (ll i = m; i < n; ++i)
#define FORR(i, m, n) for (ll i = n - 1; i >= m; --i)
#define FORE(x, xs) for (const auto &x : xs)
#define FORI(i, v) for (auto i = v.begin(); i != v.end(); i++)
#define ALL(v) v.begin(), v.end()
#define CHMIN(x, y) x = min(x, y)
#define CHMAX(x, y) x = max(x, y)
#define YES(b) cout << ((b) ? "YES" : "NO") << endl
#define Yes(b) cout << ((b) ? "Yes" : "No") << endl
#define DOUBLE(d) cout << fixed << setprecision(15) << (d) << endl

struct UnionFind {
private:
  int N;
  vi parent;
  vi rank_;

  void init(int n) {
    N = n;
    parent.resize(0);
    parent.clear();
    parent.resize(n);
    REP(i, N) parent[i] = i;
    rank_.resize(0);
    rank_.clear();
    rank_.resize(n);
  }

public:
  UnionFind() { init(0); }
  UnionFind(int n) { init(n); }

  int root(int x) { return (parent[x] == x) ? x : parent[x] = root(parent[x]); }

  bool isSame(int x, int y) { return root(x) == root(y); }

  void merge(int x, int y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return;

    if (rank_[x] < rank_[y])
      parent[x] = y;
    else {
      parent[y] = x;
      if (rank_[x] == rank_[y])
        rank_[x]++;
    }
  }
};

class BipartiteGraph {
private:
  int N;
  UnionFind uf;
  bool valid;

  void init(int n) {
    N = n;
    uf = UnionFind(2 * N);
    valid = true;
  }

public:
  BipartiteGraph() { init(0); }
  BipartiteGraph(int n) { init(n); }

  void add(int u, int v) {
    uf.merge(u, v + N);
    uf.merge(u + N, v);

    valid &= !(uf.isSame(u, u + N) || uf.isSame(v, v + N));
  }

  bool isValid() { return valid; }

  pair<set<int>, set<int>> get() {
    assert(valid);

    set<int> u, v;
    REP(i, N) {
      if (uf.isSame(i, 0))
        u.insert(i);
      else
        v.insert(i);
    }
    return make_pair(u, v);
  }
};

int N;
BipartiteGraph bi;

int main() {
  cin >> N;
  bi = BipartiteGraph(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    bi.add(a, b);
  }

  vi ans(N);
  set<int> u, v;
  tie(u, v) = bi.get();

  if (v.size() > u.size())
    swap(u, v);
  int n = 1;
  FORI(iter, u) {
    ans[*iter] = n;
    n += 3;
    if (n > N)
      break;
  }
  n = 2;
  if (v.size() >= (N + 1) / 3) {
    FORI(iter, v) {
      ans[*iter] = n;
      n += 3;
      if (n > N)
        break;
    }
  } else {
    for (auto iter = next(u.begin(), (N + 2) / 3);; iter++) {
      ans[*iter] = n;
      n += 3;
      if (n > N)
        break;
    }
  }

  n = 3;
  REP(i, N) {
    if (ans[i] != 0)
      continue;
    ans[i] = n;
    n += 3;
  }

  REP(i, N) cout << ans[i] << " ";
  cout << endl;
}