HackerEarth Optimal Edge Weights problem solution

In this HackerEarth Optimal Edge Weights problem solution There is a state (territory) which is a tree rooted at node 1. All the cities (numbered from 1 to N + 1) in this state are connected via bidirectional roads. You have to add toll tax on each road. There are N roads which connect the cities of the state. You have to assign toll taxes on the roads so as to maximize the function Toll described below:

for(i=1;i<=number of cities;i++)

{

    for(j=i+1;j<=number of cities;j++)

    {

        toll+=(toll required to pay to travel from i to j)

     } 

}

You have to maximize the toll tax. Assign the roads by toll taxes from the given array A(using each value exactly once). Find the maximum toll obtained.

HackerEarth Optimal Edge Weights problem solution.

#include<bits/stdc++.h>
using namespace std;
#define ll long long int
vector<ll>tree[100005];
ll n;
vector<pair<ll,ll> >edglist;
vector<ll>v,v1;
ll vis[100005],subtree[100005],par[100005];
void dfs(ll node)
{
    ll i,j;
    vis[node]=1;
    subtree[node]=1;
    for(i=0;i<tree[node].size();i++)
    {
        j=tree[node][i];
        if(!vis[j])
        {
            par[j]=node;
            dfs(j);
            subtree[node]+=subtree[j];
        }
    }
}
int main()
{
    freopen("samp.txt","r",stdin);
    freopen("sout.txt","w",stdout);
    ll i,j,k,x,y;
    cin>>n>>j;
    edglist.clear();
    for(i=1;i<=n;i++)
    {
        cin>>x>>y;
        tree[x].push_back(y);
        tree[y].push_back(x);
        edglist.push_back({x,y});
    }
    dfs(1);
    for(i=1;i<=n;i++)
    {
        cin>>j;
        v.push_back(j);
    }
    sort(v.begin(),v.end());
    for(i=0;i<n;i++)
    {
        x=edglist[i].first;
        y=edglist[i].second;
        if(par[x]==y)
        {
            v1.push_back(subtree[x]*(n+1-subtree[x]));
        }
        else
        {
            v1.push_back(subtree[y]*(n+1-subtree[y]));
        }
    }
    sort(v1.begin(),v1.end());
    ll ans=0;
    for(i=0;i<v.size();i++)
    {
        ans+=(v[i]*v1[i]);
    }
    cout<<ans<<"\n";
    return 0;
}

Second solution

#include<bits/stdc++.h>
#define ll long long
using namespace std;
bool vis[200005];
ll sub[200005];
vector<int>v[200005];
vector<ll>vec;
ll ans;
void dfs(int node)
{
    vis[node]=1;sub[node]++;
    for(auto u: v[node])
    {
        if(!vis[u])
        {
            dfs(u);
            sub[node]+=sub[u];
        }
    }
}
int main()
{
    ll n,two;
    cin>>n>>two;
    for(int i=1;i<=n;i++)
    {
        int x,y;
        cin>>x>>y;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    ll w[200005];
    for(int i=0;i<n;i++)
        cin>>w[i];
    dfs(1);
    sort(w,w+n);
    for(int i=2;i<=n+1;i++)
        vec.push_back((n-sub[i]+1)*sub[i]);
    sort(vec.begin(),vec.end());
    for(int i=0;i<n;i++)
        ans+=vec[i]*w[i];
    cout<<ans<<"\n";
    return 0;
}