SRM 665 Med Div2 LuckyCycle

問題文

TopCoder Statistics - Problem Statement

n個のノードとn-1個のedge(無向辺)からなる木がある（閉路がない）。
edgeには重さが4か7どちらかである。
たどっていったedgeの重さが4の個数と7の個数が一致する閉路を作れるようにふたつのノード間で繋がっていないエッジを追加しなさい(作れない場合は追加しない)。

解法

深さ優先探索で4の個数と7の個数を数えていき|4の個数-7の個数｜=1になる時のノードを根(スタート地点)とつないで閉路を作ればいい。
注意する事は深さ優先探索で全てのノードがスタート地点になる場合をためす事が必要である。

#include <iostream>
#include <vector>
#include <stdlib.h>

using namespace std;

#define WHITE 0
#define GRAY 1
#define BLACK 2

class LuckyCycle{

 public:
  vector<pair<int,int> > PQ[150];
  vector<int> answer;
  int seven_count;
  int color[150];
  int four_count;
  int need_cost;
  int start;
  int goal;
  
  void dfs(int u){
   pair<int,int>f;
   color[u]=GRAY;

   if((seven_count>1 || four_count>1) && abs(seven_count - four_count)==1
   && four_count > seven_count){
    goal=u;
    need_cost=7;
   }
   
   if((seven_count>1 || four_count>1) && abs(seven_count - four_count)==1
   && seven_count > four_count){
    goal=u;
    need_cost=4;
   }
   
   
   for(int i=0; i<PQ[u].size(); i++){
    f=PQ[u][i];
    
    if(color[f.first]==WHITE){
    
     if(f.second==4)four_count++ ;
     else if(f.second==7)seven_count++ ;

     dfs(f.first);

     if(f.second==4)four_count--;
     else if(f.second==7)seven_count--;
    }

   }

   color[u]=BLACK;   
  }

  

  vector<int>getEdge(vector <int> edge1,vector<int> edge2,vector<int> weight){
  
   for(int i=0; i<edge1.size(); i++){
     PQ[edge1[i]].push_back(make_pair(edge2[i],weight[i]));
     PQ[edge2[i]].push_back(make_pair(edge1[i],weight[i]));  
 
   }

   for(int i=0; i<=edge1.size(); i++){
    seven_count=0;
    four_count=0;
    need_cost=0;
    goal=-1;
    start=i+1;
    for(int j=0; j<=edge1.size(); j++)color[j+1]=WHITE;

    dfs(start);
    if(goal!=-1 && need_cost!=0)break;
   }

   if(goal!=-1 && need_cost!=0){
    answer.push_back(start);
    answer.push_back(goal);
    answer.push_back(need_cost);
   }

   return answer;  
  }
};