SRM559 Div1 Medium(500) HatRack
ある部分木が正しいラックになる条件は、子の個数が2個以下で、両方の子の高さが同じで少なくとも一方が完全二分木であること、もしくは両方の子の高さの差が1で、短い方が完全部分木であること。子が1個しかなければ、もう一方は高さが0の完全二分木と見なす。また、子の高さが同じでどちらも完全二分木ならば、子は交換することができる。
#include <vector> #include <algorithm> using namespace std; vector<vector<int> > E; // ノードpを親としてノードcを辿ったとき、何通りのラックができるかを返す // dに木の高さ、fに木が完全二分木かどうかを返す long long BT( int c, int p, int *d, bool *f ) { vector<int> e = E[c]; e.erase(remove(e.begin(),e.end(),p),e.end()); if ( (int)e.size()==0 ) { *d = 1; *f = true; return 1LL; } if ( (int)e.size() > 2 ) return 0LL; long long r1, r2; int d1, d2; bool f1, f2; r1 = BT( e[0], c, &d1, &f1 ); if ( (int)e.size()==2 ) { r2 = BT( e[1], c, &d2, &f2 ); } else { r2 = 1LL; d2 = 0; f2 = true; } if ( d1==d2 && ( f1 || f2 ) ) { *d = d1 + 1; *f = f1 && f2; return ( f1&&f2 ? 2 : 1 ) * r1*r2; } if ( d1==d2+1 && f2 || d2==d1+1 && f1 ) { *d = max(d1,d2)+1; *f = false; return r1*r2; } return 0LL; } class HatRack{public: long long countWays( vector <int> knob1, vector <int> knob2 ) { int N = (int)knob1.size() + 1; E = vector<vector<int> >( N ); for ( int i=0; i<N-1; i++ ) E[knob1[i]-1].push_back(knob2[i]-1), E[knob2[i]-1].push_back(knob1[i]-1); long long ans = 0; int d; bool f; for ( int i=0; i<N; i++ ) ans += BT( i, -1, &d, &f ); return ans; }};