喵哈哈村的魔法考试 Round #10 (Div.2) D

喵哈哈村与哗啦啦村的大战（四）

发布时间: 2017年3月27日 10:10   时间限制: 1000ms   内存限制: 128M

描述

喵哈哈村因为和哗啦啦村争夺稀有的水晶资源，展开了激烈的战斗！

喵哈哈村的部落可以视为由n个节点组成，其中有n-1条边连接这n个节点，使得任意两个节点都会有一条路径相连接。每个节点上都有一个点权a[i]。

如果说存在一条路径上的权值满足非严格的单调递增或者非严格的单调递减的话，就说这条路径是一条好路径。

现在问题来了，给你一棵树，问你这棵树上有多少条路径是好路径。

输入

本题包含若干组测试数据。
第一行一个n，表示有n个节点。
第二行n个整数a[i]，表示节点的权值。
接下来n-1行每行两个整数x,y。表示x,y节点之间有一条边相互连接。

满足：1<=n<=100000 1<=a[i]<=1e9 1<=x,y<=n

输出

输出好路径的个数。

样例输入1 复制

4
1 7 1 9
1 3
1 4 
2 1

样例输出1

5

样例输入2 复制

6
1 1 2 2 3 3
1 2
2 3
3 4
4 5
5 6

样例输出2

15

题解：树形dp，因为是树，随便找个点作为根就可以，然后往下面搜，将叶子信息往上传更新父节点的信息就好了。

#include<bits/stdc++.h>
using namespace std;
#define PB push_back
typedef  long long ll;
typedef  pair<int, int> PII;

const int maxn=1e6+5;
vector<int> G[maxn];
int a[maxn], dp[maxn][3], ans;

void dfs(int x, int fa){
    for(int i=0; i<G[x].size(); i++){
        int v=G[x][i];
        if(v==fa)    continue;
        dfs(v, x);
        if(a[x]<a[v]){
            ans+=(dp[x][1]+dp[x][0])*(dp[v][2]+dp[v][0]+1);
            dp[x][2]+=dp[v][2]+dp[v][0]+1;
        }else if(a[x]>a[v]){
            ans+=(dp[x][2]+dp[x][0])*(dp[v][1]+dp[v][0]+1);
            dp[x][1]+=dp[v][1]+dp[v][0]+1; 
        }else{
            ans+=dp[x][0]*(dp[v][1]+dp[v][2]+dp[v][0]+1)+dp[x][1]*(dp[v][2]+dp[v][0]+1)+dp[x][2]*(dp[v][1]+dp[v][0]+1);
            dp[x][0]+=dp[v][0]+1;
            dp[x][1]+=dp[v][1];
            dp[x][2]+=dp[v][2];
        }
    }
    ans+=dp[x][0]+dp[x][1]+dp[x][2];
}

int main(){
    int n;
    while(cin>>n){
        memset(dp, 0, sizeof(dp));
        for(int i=1; i<=n; i++)    G[i].clear();
        for(int i=1; i<=n; i++){
            cin>>a[i];
        }
        int a, b;
        for(int i=1; i<=n-1; i++){
            cin>>a>>b;
            G[a].PB(b), G[b].PB(a);
        }
        ans=0;
        dfs(1, 0);
        cout<<ans<<endl;
    }
    return 0;
}