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