 # Editorial: PRT

Practice

Contest: Division 1

Contest: Division 2

Editorialist: Raja Vardhan Reddy

Easy

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;
int sub[M], leaves, mod = 1e9 + 7;
long long paths[M];

int calc(int cur, int parent) {
int r = 0;

if (x != parent)
r += calc(x, cur);

sub[cur] = r;
return r;
}

void solve(int cur, int parent) {
int sum = 0;
vector<int> p;

if (x != parent)
p.push_back(sub[x]), sum += sub[x], solve(x, cur);
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++)
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--;
}

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

return [int(x) for x in input().split()]

st = [src]
ret = []
visited = [False] * (n + 1)
visited[src] = True

while len(st) > 0:
u = st.pop()
ret.append(u)
if not visited[v]:
st.append(v)
visited[v] = True

return ret

# -----------------#
#       Main      #
# -----------------#

for test in range(T):

adj = [[] for i in range(n + 1)]

for i in range(n - 1):

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):
root = i

sz =  * (n + 1)
par = [-1] * (n + 1)

cnt_leafs = 0
for u in dfs_order[::-1]:
sz[u] += sz[v]
if sz[v] == 0:
par[u] = v

sz[u] += 1
cnt_leafs += 1

coef =  * (n + 1)
for u in range(1, n + 1):
s = 0
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 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,l,deg,a;
int leaf=0;
int dfs(int u,int p){
if(deg[u]==1){
l[u]=1;
}
else{
l[u]=0;
}
int i,s=0;
}
}
int val=0;
}
}
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;
deg[i]=0;
}
for(i=0;i<n-1;i++){
cin>>u>>v;
u--;
v--;
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. 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 ?

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: 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.