Editorial: PRT

PROBLEM LINK:

Practice

Contest: Division 1

Contest: Division 2

Setter: Ahmad Salah

Tester: Radoslav Dimitrov

Editorialist: Raja Vardhan Reddy

DIFFICULTY:

Easy

PREREQUISITES:

dfs, greedy

PROBLEM:

You are given a tree with N vertices (numbered 1 through N) and a sequence of integers A_1, A_2, \cdots A_N. You may choose an arbitrary permutation p_1, p_2, \cdots p_N of the integers 1 through N. Then, for each vertex i, you should assign the value A_{p_i} to this vertex.

The profit of a path between two vertices u and v is the sum of the values assigned to the vertices on that path (including u and v).

Let’s consider only (undirected) paths that start at a leaf and end at a different leaf. Calculate the maximum possible value of the sum of profits of all such paths. Since this value could be very large, compute it modulo 10^9+7.

EXPLANATION

Let A_{p_i} be the value assigned to vertex i and let vertex i occur on c_i different paths. Then, the sum of profits over all paths = Σ(c_i*A_{p_i}) over all the vertices. For maximizing the profit, it is optimal to assign largest A_i to the vertex which occurs on more paths.So, let us calculate c_i for each i, and greedily assign largest A_i to the vertex with maximum c_i.

Calculation of c_i for a vertex i:

Root the tree at a non-leaf node. Now, vertex i will be present on:
a.) All the paths which have one leaf in the subtree of vertex i and the other one which is not in the subtree of vertex i.
b.) All the paths which have one leaf in one of its child’s subtree and the other leaf in its any other child’s subtree.

Let l be the total number of leaves in the tree, and leaf_i be the number if leaves in subtree of vertex i.
Number of paths of type a: leaf_i*(l-leaf_i);
Number of paths of type b: Σ(leaf_j*(leaf_i-leaf_j))/2 where j children of i.

Hence, c_i=leaf_i*(l-leaf_i)+ (Σ(leaf_j*(leaf_i-leaf_j))/2 over all children of vertex i).
l and leaf_i can be calculated by running a dfs on the rooted tree.

After calculation of c_i for all vertices, sort the vertices in the decreasing order of c_i,sort the given array A in decreasing order, and calculate s=Σ(c_i*A_{i}) , which is the required answer.

TIME COMPLEXITY:

Calculation of n_i for all vertices : O(n)
Sorting the verices w.r.t c_i : O(nlog(n))
Sorting the array A : O(nlog(n))
Overall complexity: O(nlog(n)+nlog(n)+n) = O(nlog(n)) for each test case.

SOLUTIONS:

Setter's Solution
#include <bits/stdc++.h>
 
using namespace std;
 
const int M = 3e5;
vector<int> adj[M];
int sub[M], leaves, mod = 1e9 + 7;
long long paths[M];
 
int calc(int cur, int parent) {
    int r = 0;
 
    for (int x : adj[cur])
        if (x != parent)
            r += calc(x, cur);
 
    r += (adj[cur].size() == 1);
    leaves += (adj[cur].size() == 1);
 
    sub[cur] = r;
    return r;
}
 
void solve(int cur, int parent) {
    int sum = 0;
    vector<int> p;
 
    for (int x : adj[cur])
        if (x != parent)
            p.push_back(sub[x]), sum += sub[x], solve(x, cur);
    if (adj[cur].size() == 1)
        p.push_back(1), sum++;
 
    p.push_back(leaves - sum);
    sum = leaves;
 
    for (int x : p)
        sum -= x, paths[cur] += 1ll * x * sum;
}
 
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
 
    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        int arr[n], ans = 0;
 
        for (int i = 0; i < n; i++)
            adj[i].clear();
        memset(paths, 0, sizeof paths);
        leaves = 0;
 
        for (int i = 0; i < n; i++)
            cin >> arr[i];
 
        for (int i = 0; i < n - 1; i++) {
            int x, y;
            cin >> x >> y;
            x--;
            y--;
            adj[x].push_back(y);
            adj[y].push_back(x);
        }
 
        calc(0, -1);
        solve(0, -1);
 
        sort(paths, paths + n);
        sort(arr, arr + n);
 
        for (int i = 0; i < n; i++)
            ans += (1ll * (paths[i] % mod) * arr[i]) % mod, ans %= mod;
 
        cout << ans << '\n';
    }
 
    return 0;
} 
Tester's Solution
inf = 10 ** 18
mod = 10 ** 9 + 7
 
def read_line_int():
    return [int(x) for x in input().split()]
 
def getDfsOrder(adj, src):
    st = [src]
    ret = []
    n = len(adj) - 1
    visited = [False] * (n + 1)
    visited[src] = True
 
    while len(st) > 0:
        u = st.pop()
        ret.append(u)
        for v in adj[u]:
            if not visited[v]:
                st.append(v)
                visited[v] = True
 
    return ret
 
# -----------------#
#       Main      #
# -----------------#
 
 
 
T = read_line_int()[0]
 
for test in range(T):
    n = read_line_int()[0]
    a = read_line_int()
 
    adj = [[] for i in range(n + 1)]
 
    for i in range(n - 1):
        u, v = read_line_int()
        adj[u].append(v)
        adj[v].append(u)
    
    if n == 1:
        print(0)
        continue
 
    if n == 2:
        print(sum(a) % mod)
        continue
 
    mx_deg = -1
    root = -1
    for i in range(1, n + 1):
        if mx_deg < len(adj[i]):
            mx_deg = len(adj[i])
            root = i
 
    dfs_order = getDfsOrder(adj, root)
    sz = [0] * (n + 1)
    par = [-1] * (n + 1)
 
    cnt_leafs = 0
    for u in dfs_order[::-1]:
        for v in adj[u]:
            sz[u] += sz[v]
            if sz[v] == 0:
                par[u] = v
 
        if len(adj[u]) == 1:
            sz[u] += 1
            cnt_leafs += 1
 
    coef = [0] * (n + 1)
    for u in range(1, n + 1):
        s = 0
        for v in adj[u]:
            if v == par[u]: continue
            coef[u] += s * sz[v]
            s += sz[v]
 
        coef[u] += sz[u] * (cnt_leafs - sz[u])
 
    coef = coef[1:]
    coef = sorted(coef)
    a = sorted(a)
 
    ret = sum([x * y % mod for (x, y) in zip(coef, a)]) % mod
    print(ret)
Editorialist's Solution
    //raja1999

//#pragma comment(linker, "/stack:200000000")
//#pragma GCC optimize("Ofast")
//#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")

#include <bits/stdc++.h>
#include <vector>
#include <set>
#include <map>
#include <string> 
#include <cstdio>
#include <cstdlib>
#include <climits>
#include <utility>
#include <algorithm>
#include <cmath>
#include <queue>
#include <stack>
#include <iomanip> 
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp> 
//setbase - cout << setbase (16)a; cout << 100 << endl; Prints 64
//setfill -   cout << setfill ('x') << setw (5); cout << 77 <<endl;prints xxx77
//setprecision - cout << setprecision (14) << f << endl; Prints x.xxxx
//cout.precision(x)  cout<<fixed<<val;  // prints x digits after decimal in val

using namespace std;
using namespace __gnu_pbds;
#define f(i,a,b) for(i=a;i<b;i++)
#define rep(i,n) f(i,0,n)
#define fd(i,a,b) for(i=a;i>=b;i--)
#define pb push_back
#define mp make_pair
#define vi vector< int >
#define vl vector< ll >
#define ss second
#define ff first
#define ll long long
#define pii pair< int,int >
#define pll pair< ll,ll >
#define sz(a) a.size()
#define inf (1000*1000*1000+5)
#define all(a) a.begin(),a.end()
#define tri pair<int,pii>
#define vii vector<pii>
#define vll vector<pll>
#define viii vector<tri>
#define mod (1000*1000*1000+7)
#define pqueue priority_queue< int >
#define pdqueue priority_queue< int,vi ,greater< int > >
#define int ll

typedef tree<
int,
null_type,
less<int>,
rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;


//std::ios::sync_with_stdio(false);

int arr[300005],l[300005],deg[300005],a[300005];
vector<vector<int> >adj(300005);
int leaf=0;
int dfs(int u,int p){
	if(deg[u]==1){
		l[u]=1;
	}
	else{
		l[u]=0;
	}
	int i,s=0;
	for(i=0;i<adj[u].size();i++){
		if(adj[u][i]!=p){
			dfs(adj[u][i],u);
			l[u]+=l[adj[u][i]];
			s+=l[adj[u][i]];
		}
	}
	int val=0;
	for(i=0;i<adj[u].size();i++){
		if(adj[u][i]!=p){
			val+=(s-l[adj[u][i]])*(l[adj[u][i]]);
		}
	}
	val/=2;
	val+=l[u]*(leaf-l[u]);
	arr[u]=val;
	return 0;
}
main(){
	std::ios::sync_with_stdio(false); cin.tie(NULL);
	int t,t1;
	cin>>t;
	//t=1;
	t1=t;
	while(t--){
		int n,i,sum=0,u,v,root,ans=0;
		cin>>n;
		leaf=0;
		for(i=0;i<n;i++){
			cin>>a[i];
			sum+=a[i];
			sum%=mod;
			adj[i].clear();
			deg[i]=0;
		}
		for(i=0;i<n-1;i++){
			cin>>u>>v;
			u--;
			v--;
			adj[u].pb(v);
			adj[v].pb(u);
			deg[u]++;
			deg[v]++;
		}
		if(n==1){
			cout<<0<<endl;
			continue;
		}
		if(n==2){
			cout<<sum<<endl;
			continue;
		}
		for(i=0;i<n;i++){
			if(deg[i]==1){
				leaf++;
			}
			else{
				root=i;
			}
		}
		dfs(root,-1);
		sort(arr,arr+n);
		sort(a,a+n);
		for(i=0;i<n;i++){
			arr[i]%=mod;
			ans+=(arr[i]*a[i])%mod;
			ans%=mod;
		}
		cout<<ans<<endl;
	}
	return 0;
} 

Feel free to Share your approach, If it differs. Suggestions are always welcomed. :slight_smile:

7 Likes

i have used the same approach which u mentioned above. but i m getting WA?
will u pls tell me what i am missing ?
my submission link : https://www.codechef.com/viewsolution/29091987

I used exact same approach but still my solution fails. Any counter case?
https://www.codechef.com/viewsolution/29094630

You are applying mod while calculating number of paths in which the node occurs, which is wrong. You need to calculate it without mod.

1 Like

Oh ya, it create problem for sorting. Thank you very much.

1 Like

yes. you are right …
thank u so much bridarrrrrr: :heart_eyes:

i am also getting wrong ans but approch same help me.
https://www.codechef.com/viewsolution/29099228

My logic is same but i am getting wrong answer.
my submission: https://www.codechef.com/viewsolution/29102761

Instead of vertices problem can be constructed for edges.how to approach it then?

1 Like

I think number of paths of type b has to be divided by 2 while calculating ci, since it is being calculated twice. I see in the Editorialist’s code it is done so. It would be good if the post is edited too.

1 Like

I’m not able to understand why we’re only considering leaves for the occurance for given node? We should be considering no. of nodes in it’s subtree.

fixed. Thanks.

1 Like

Read the question again. We are concerned about paths between two different leaves only.

Okay I got it. It’s given in 3rd paragraph. Forgot to read whole question during contest, did same mistake while reading article.

I am using dfs, still getting TLE. Please help me out.

https://www.codechef.com/viewsolution/29134129

i am not able to understand how we have calculated ci for each vertex i . someone either please explain me or if it is a standard approach please suggest me the topic and i will have a look on tutorial …please help

You have to find the number of paths that passes through a vertex, that is what meant by C_{i}. Now, in order for a path to pass through the vertex i, at-least one leaf node must lie in it’s sub-tree(why?). So, you can have either both the leaves lying in the subtree of the vertex or only one. For the case when both the leaves are in the subtree, if you have N_{V} leaves in the subtree of the vertex V_{i}, then you have to choose two leaves to construct a unique path, this can be done by - \binom{N_{V}}{2}. Now, if total number of leaves in the tree is N_{T} then the leaves that are not in the sub-tree of V are N_{T} - N_{V}. So, to form a path passing from V, such that only one leaf is in sub-tree of V, the answer is (N_{T} - N_{V}) \ * \ N_{V}. So, sum up both, i.e. C_{i} = (N_{T} - N_{V}) \ * \ N_{V} + \binom{N_{V}}{2}.

1 Like