1092F - Tree with Maximum Cost - CodeForces Solution


dfs and similar dp trees *1900

Please click on ads to support us..

Python Code:

from math import inf
from collections import *
import math, os, sys, heapq, bisect, random,threading
from functools import lru_cache
from itertools import *
import sys
def inp(): return sys.stdin.readline().rstrip("\r\n")
def out(var): sys.stdout.write(str(var))  def inpu(): return int(inp())
def lis(): return list(map(int, inp().split()))
def stringlis(): return list(map(str, inp().split()))
def sep(): return map(int, inp().split())
def strsep(): return map(str, inp().split())
def fsep(): return map(float, inp().split())
M,M1=1000000007,998244353


def main():
    how_much_noob_I_am = 1
        for _ in range(1,how_much_noob_I_am+1):
        n = inpu()
        arr = lis()
        d=defaultdict(list)
        for i in range(n-1):
            a,b = sep()
            d[a].append(b)
            d[b].append(a)
        vis = [False]*(n+1)
        ans = [0]*(n+1)
        res = 0
        def dfs(curr,depth):
            nonlocal res
            res+=depth*arr[curr-1]
            ans[curr]=arr[curr-1]
            vis[curr] = True
            for j in d[curr]:
                if not vis[j]:
                    dfs(j,depth+1)
                    ans[curr]+=ans[j]
        dfs(1,0)
        result = 0
        def dfs2(curr,par):
            nonlocal ans,result,res
            result = max(res,result)
            for j in d[curr]:
                if j!=par:
                    res-=ans[j]
                    ans[curr]-=ans[j]
                    res+=ans[curr]
                    ans[j]+=ans[curr]

                    dfs2(j,curr)

                    ans[j] -= ans[curr]
                    res -= ans[curr]
                    ans[curr] += ans[j]
                    res+=ans[j]



        dfs2(1,-1)
        print(result)



if __name__ == '__main__':
    sys.setrecursionlimit(2*10**5+50)
    threading.stack_size(10**8)
    threading.Thread(target=main).start()
    

C++ Code:

#include <bits/stdc++.h>
using namespace std;
#define ff first
#define ss second
#define f0(i, n) for (int i = 0; i < n; i++)
#define f1(i, n) for (int i = 1; i <= n; i++)
#define int long long
#define all(x) x.begin(), x.end()
#define pb push_back
#define pii pair<int, int>
#define vi vector<int>
#define mii map<int, int>
#define mod7 1000000007
#define t_case int t; cin>>t; while(t--)
#define setp(x, y) fixed << setprecision(y) << x
#define setbits(x) __builtin_popcountll(x)
#define zerobits(x) __builtin_ctzll(x)
#define fast_io std::ios::sync_with_stdio(false), cin.tie(NULL), cout.tie(NULL)
const int N=2e5+10;
int sum[N],cost[N],a[N];
int n;
int ans;
vector<int>ad[N];
void dfs1(int u=1,int p=-1){
	sum[u]=a[u];
	for(auto i:ad[u]){
		if(i!=p){
			dfs1(i,u);
			sum[u]+=sum[i];
			cost[u]+=cost[i]+sum[i];
		}
	}
}
void dfs2(int u=1,int p=-1){
	ans=max(ans,cost[u]);
	for(auto i:ad[u]){
		if(i!=p){
			int newcost=cost[u]-(sum[i]+cost[i]);
			cost[i]+=newcost+sum[u]-sum[i];
			sum[i]+=sum[u]-sum[i];
			dfs2(i,u);
		}
	}
}
void solve(){
cin>>n;
f1(i,n) cin>>a[i];
f1(i,n-1){
	int u,v;cin>>u>>v;
	ad[u].pb(v);
	ad[v].pb(u);
}
dfs1();
dfs2();
cout<<ans<<endl;


}
signed main()
{
    fast_io;
    // t_case{
    	solve();
    // }



}


Comments

Submit
0 Comments
More Questions

1302. Deepest Leaves Sum
1209. Remove All Adjacent Duplicates in String II
994. Rotting Oranges
983. Minimum Cost For Tickets
973. K Closest Points to Origin
969. Pancake Sorting
967. Numbers With Same Consecutive Differences
957. Prison Cells After N Days
946. Validate Stack Sequences
921. Minimum Add to Make Parentheses Valid
881. Boats to Save People
497. Random Point in Non-overlapping Rectangles
528. Random Pick with Weight
470. Implement Rand10() Using Rand7()
866. Prime Palindrome
1516A - Tit for Tat
622. Design Circular Queue
814. Binary Tree Pruning
791. Custom Sort String
787. Cheapest Flights Within K Stops
779. K-th Symbol in Grammar
701. Insert into a Binary Search Tree
429. N-ary Tree Level Order Traversal
739. Daily Temperatures
647. Palindromic Substrings
583. Delete Operation for Two Strings
518. Coin Change 2
516. Longest Palindromic Subsequence
468. Validate IP Address
450. Delete Node in a BST