【DFS】HDU 5423 Rikka with Tree

/* 
DFS
Rikka with Tree
時間: 2017/02/24
題意：判斷樹是否存在跟他相似並不同的樹
題解： 
要求相似並不同，樹的結構必須除了最後一層，每層只能一個點
最開始我用記憶化搜尋去計算其長度，因為我覺得樹都是已知指向未知，相當於DAG，但是一直wa。（困惑）
後來用dfs計算長度過了
*/  

#include<cstdio>
#include<cstring>
#include<cmath>
#include<iostream>
#include<algorithm>
#include<queue>
#include<map>
using namespace std;
#define LL long long
#define N 1010
#define INF 0x3f3f3f3f

int mp[N][N],dist[N],vis[N];
int n;
int DFS(int i)
{
    if(dist[i] >= 0)
        return dist[i];
    for(int j = 1; j < n; j++)
    {
        if(mp[i][j] && dist[j] >= 0)
            return dist[i] = 1+DFS(j);
    }
}
void dfs(int u,int depth)
{
    dist[u]=depth;
    //printf("%d   %d\n",u,dist[u]);
    for(int i = 1; i <= n; i++)
    {
        vis[u]=true;
        if(mp[i][u] && !vis[i])
        {
            dfs(i,depth+1);
            vis[u]=false;
        }
    }
}
int main()
{
    while(~scanf("%d",&n))
    {
        memset(mp,0,sizeof(mp));
        int a,b;
        for(int i = 1; i < n; i++)
        {
            scanf("%d%d",&a,&b);
            mp[a][b] = mp[b][a] = 1;
        }
        int flag = 1;
        memset(dist,-1,sizeof(dist));
        memset(vis,0,sizeof(vis));
        int mx = -1;
        dfs(1,0);
        for(int i = 1; i <= n; i++)
        {
            mx = max(mx,dist[i]);
        }
        for(int i = 2; i <= n; i++)
        {
            for(int j = i+1; j <= n; j++)
                if(dist[i] == dist[j] && dist[i] != mx)
                {
                    flag = 0;
                    break;
                }
            if(!flag)
                break;
        }
        if(flag)
            puts("YES");
        else
            puts("NO");
    }
    return 0;
}