SRM 644 Div2 1000 TreeCutting

問題概要
木がある。
木のノードのいくつかには数字が書かれている。
木を以下の条件を満たすように、いくつかのコンポーネントに分割することは可能か？

• 各コンポーネントの中に数字が書かれたノードは一つだけ存在する
• 各コンポーネント内のノード数と、コンポーネント内のノードに書かれた数字は等しい

解法
枝(v, w)で木を分断するとする。
このとき、v側のコンポーネントとw側のコンポーネントに分かれる。
各コンポーネント内のノード数と、各コンポーネント内の数値の合計が等しければ、そこで分断してもいいことがわかる。
このような分断の個数は、解状態におけるコンポーネント数 - 1個だけあるはず。
よって、分断数 + 1 = 数字が書かれたノードの数をチェックすればいい。

実装

using namespace std;

#define ALL(x) (x).begin(), (x).end()
#define EACH(itr,c) for(__typeof((c).begin()) itr=(c).begin(); itr!=(c).end(); itr++)  
#define FOR(i,b,e) for (int i=(int)(b); i<(int)(e); i++)
#define MP(x,y) make_pair(x,y)
#define REP(i,n) for(int i=0; i<(int)(n); i++)

vector<int> nums;
vector<int> g[50];

class TreeCutting {
  int cnt(int v, int par) {
    int ret = 1;
    for (auto &w: g[v])
      if (w != par)
        ret += cnt(w, v);
    return ret;
  }

  int sum(int v, int par) {
    int ret = max(0, nums[v]);
    for (auto &w: g[v])
      if (w != par)
        ret += sum(w, v);
    return ret;
  }

public:
  string can(vector<int> par, vector<int> num) {
    int n = num.size();
    int m = 0;
    nums = num;
    REP (i, n) if (num[i] != -1) ++m;
    REP (i, n) g[i].clear();
    REP (i, par.size()) {
      g[par[i]].push_back(i+1);
      g[i+1].push_back(par[i]);
    }
    
    int split = 0;
    REP (i, par.size()) {
      int v = par[i];
      int w = i+1;
      if (cnt(v, w) == sum(v, w) && cnt(w, v) == sum(w, v))
        ++split;
    }

    if (split+1 == m)
      return "POSSIBLE";
    else
      return "IMPOSSIBLE";
  }
};