Topcoder SRM 633 Div. 1 600 - DoubleTree

Problem Statement:
SRM 633 Div. 1 600 - DoubleTree

Solution:
It is interesting to know that this problem is actually a Maximum Flow problem. When we fix a node as the root of both of the trees, we reduce the problem to maximum weight connected sub graph problem. Let's say we have chosen a root R. R will be inside the subset S in question. We run a depth first search on both tree starting at R, forming two trees A and B. Then we build a graph G such that:
1. for each node v, we connect v with its parent u in A and B each with an edge with infinite capacity. This means that if we pick node v, then we have to also pick u, so we ensure that v and R is connected.
2. introduce two new nodes, source s and sink t. For each v in G, if v has a positive score m, add an edge between s to v with capacity m. Otherwise, m is negative, then we connect v with t using an edge with capacity -m.

3. Furthermore, we sum up all positive scores as SUM. SUM represents the maximum score possible if it is possible to choose all positive nodes.
4. Run maxflow algorithm, and we will get the maximum flow MF, which is also the minimum cut between s and t.
5. SUM - MF represents the maximum weight S if R is in S. Why? Whenever we choose a positive node v that is not directly connected to R, we have to go through its parents. If the parents have a negative score, then s will send flow to t by that negative score. Hence in the end, MF will represent the amount of negative scores incurred as a result of choosing all the positive nodes. Hence SUM - MF will give us the total weight of S. (cool stuff!)

Implementation:

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
#include <utility>
#include <queue>
using namespace std;

int INF = 1234567;
int S = 55;
int T = 56;
vector<vector<int> > adj[2];
vector<vector<int> > G;
int cap[60][60];
int flow[60][60];
int par[60];
int vis[2][60];
int N;

void clear_state(int N) {
 for(int i=0;i<N+3;++i){
  for(int j=0;j<N+3;++j){
   cap[i][j] = flow[i][j] = 0;
  }
 }
 for(int i=0;i<2;++i){
  for(int j=0;j<N;++j){
   vis[i][j] = 0;
  }
 }
 G = vector<vector<int> > (N+3);
}

void dfs(int t, int u) {
 vis[t][u] = 1;
 for(int i=0;i<adj[t][u].size();++i){
  int v = adj[t][u][i];
  if(vis[t][v]) continue;
  dfs(t, v);
  cap[v][u] = INF;
  G[u].push_back(v);
  G[v].push_back(u);
 }
}

void init(vector<int>& score) {
 for(int i=0;i<N;++i) {
  if(score[i] > 0) {
   cap[S][i] = score[i];
   G[S].push_back(i);
   G[i].push_back(S);
  } else {
   cap[i][T] = -score[i];
   G[T].push_back(i);
   G[i].push_back(T);
  }
 }
}


int maxflow;

int augment_path(int v, int f) {
 if(v == S) {
  return f;
 }
 int u = par[v];
 int backedge = 0;
 int mf = 0;
 if(cap[u][v] > 0) {
  backedge = 0;
  mf = cap[u][v] - flow[u][v];
 } else {
  backedge = 1;
  mf = flow[v][u];
 }
 int minf = augment_path(u, min(mf, f));
 if(!backedge) {
  flow[u][v] += minf;
 } else {
  flow[v][u] -= minf;
 }
 return minf;
}

void edmondskarp() {
 maxflow = 0;
 while(1) {
  for(int i=0;i<60;++i) par[i] = -1;
  queue<int> q;
  q.push(S);
  par[S] = S;
  bool found = false;
  while(!q.empty()) {
   int u = q.front();
   q.pop();
   for(int i=0;i<G[u].size(); ++i){
    int v = G[u][i];
    if(par[v] != -1) continue;
    if(cap[u][v] - flow[u][v] > 0 || flow[v][u] > 0) {
     par[v] = u;
     if(v == T) {
      found = true;
      break;
     } else {
      q.push(v);
     }
    }
   }
   if(found) break;
  }
  if(found) {
   maxflow += augment_path(T, INF);
  } else {
   break;
  }
 }
}

class DoubleTree {
public:
 int maximalScore(vector<int> a, vector<int> b, vector<int> c, vector<int> d, vector<int> score) {
  //SPECIAL CASE: when all scores are positive, just return the whole tree
  N = score.size();
  S = N+1;
  T = N+2;
  for(int i=0;i<2;++i){
   adj[i] = vector<vector<int> >(N+3);
  }
  for(int i=0;i<N-1;++i) {
   adj[0][a[i]].push_back(b[i]);
   adj[0][b[i]].push_back(a[i]);
  }
  for(int i=0;i<N-1;++i){
   adj[1][c[i]].push_back(d[i]);
   adj[1][d[i]].push_back(c[i]);
  }
  int sum = 0;
  bool special_case = true;
  for(int i=0;i<N;++i){
   if(score[i] > 0) sum += score[i];
   else special_case = false;
  }
  if(special_case || sum == 0) {
   return sum;
  } else {
   int ans = 0;
   for(int i=0;i<N;++i) {
    clear_state(N);
    dfs(0,i);
    dfs(1,i);
    init(score);
    edmondskarp(); 
    int tmp = sum - maxflow;
    ans = max(ans, tmp);
   }
   return ans;
  }
 }
};