SRM559 Div1 Medium(500) HatRack

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;
}};