Codeforces 486D. Valid Sets

486D - Valid Sets

• Firstly, we solve the case d =  + ∞. In this case, we can forget all ai since they doesn't play a role anymore. Consider the tree is rooted at node 1. Let Fi be the number of valid sets contain node i and several other nodes in subtree of i("several" here means 0 or more). We can easily calculate Fi through Fj where j is directed child node of i: . Complexity: O(n).

• General case: d ≥ 0. For each node i, we count the number of valid sets contain node i and some other node j such thatai ≤ aj ≤ ai + d (that means, node i have the smallest value a in the set). How? Start DFS from node i, visit only nodes jsuch that ai ≤ aj ≤ ai + d. Then all nodes have visited form another tree. Just apply case d =  + ∞ for this new tree. We have to count n times, each time is O(n), so the overall complexity is O(n2). (Be craeful with duplicate counting)

Here is my code.

D. Valid Sets

time limit per test

1 second

memory limit per test

256 megabytes

input

standard input

output

standard output

As you know, an undirected connected graph with n nodes and n - 1 edges is called a tree. You are given an integer d and a tree consisting of n nodes. Each node i has a value ai associated with it.

We call a set S of tree nodes valid if following conditions are satisfied:

1. S is non-empty.
2. S is connected. In other words, if nodes u and v are in S, then all nodes lying on the simple path between u and v should also be presented in S.
3. .

Your task is to count the number of valid sets. Since the result can be very large, you must print its remainder modulo1000000007 (109 + 7).

Input

The first line contains two space-separated integers d (0 ≤ d ≤ 2000) and n (1 ≤ n ≤ 2000).

The second line contains n space-separated positive integers a1, a2, ..., an(1 ≤ ai ≤ 2000).

Then the next n - 1 line each contain pair of integers u and v (1 ≤ u, v ≤ n) denoting that there is an edge between u and v. It is guaranteed that these edges form a tree.

Output

Print the number of valid sets modulo 1000000007.

Sample test(s)

input

1 4
2 1 3 2
1 2
1 3
3 4

output

8

input

0 3
1 2 3
1 2
2 3

output

3

input

4 8
7 8 7 5 4 6 4 10
1 6
1 2
5 8
1 3
3 5
6 7
3 4

output

41

Note

In the first sample, there are exactly 8 valid sets: {1}, {2}, {3}, {4}, {1, 2}, {1, 3}, {3, 4} and {1, 3, 4}. Set {1, 2, 3, 4}is not valid, because the third condition isn't satisfied. Set {1, 4} satisfies the third condition, but conflicts with the second condition.

ll dp[N];
void dfs(int u, int fa, int rt)
{
    dp[u] = 1;
    for(int i=0;i<g[u].size();i++)
    {
        int v = g[u][i];
        if(v == fa) continue;
        if(a[rt]<=a[v] && a[v]<=a[rt]+d)
        {
            if(a[v]==a[rt] && rt>v) continue;
            dfs(v, u, rt);
            dp[u] = dp[u] * (dp[v] + 1) % mod;
        }
    }
}
int main()
{
    while(cin>>d>>n)
    {
        for(int i=1;i<=n;i++)
        {
            cin>>a[i];
            g[i].clear();
        }
        for(int i=1;i<n;i++)
        {
            int u, v;
            cin>>u>>v;
            g[u].push_back(v);
            g[v].push_back(u);
        }
        ll ans = 0;
        for(int i=1;i<=n;i++)
        {
            memset(dp, 0, sizeof dp);
            dfs(i, -1, i);
            printf("root = %d\n", i);
            for(int j=1;j<=n;j++)
            {
                if(dp[j]) printf("dp[%d] = %d\n", j, dp[j]);
            }
            puts("========");
            ans = (ans + dp[i]) % mod;
        }
        cout<<ans<<endl;
    }
}