Segment Trees #4 Segment Tree on Rooted Tree: Point Queries + DP

JosephLand Problem Code: JTREE    

DP + Segment Tree

Problem Statement: here

The problem, in short, says: You have a tree with edges directed towards the root and nodes having some tickets. Each ticket allows you to move k units forward to your capital with a cost w. 

We need to answer queries which ask us for a minimum cost of traveling from node x.

This problem is tricky because of the larger constraints on the values of N, M, and Q. This is a problem of the type where you use Segment Trees as an optimization technique. 

Approach:

Let us say, we are at node x in our tree and we wish to get the minimum possible cost to travel to node 1(the capital city). This node is at depth H. So, it has H-1 nodes in its path. 

Let us break the problem into a simple version:
We want to decide which ticket is the best option among the available tickets at that node.
We will keep track of the minimum possible cost from a node in the dp[node].
We can use the recurrence for calculating the minimum possible cost: THIS IS A PSEUDO CODE.

for(ticket : all tickets){ //Iterate over all tickets available at node X cost = ticket.cost; k = ticket.k; for(nodes:nodes_in_path_of_X_upto_k_distance){ dp[X] = min(dp[X],dp[nodes] + cost); } }

Here we are iterating over all possible tickets at node X and then iterating over all the nodes up to depth k above it and taking the minimum. This makes the complexity of code O(Q + N*M)

But, we can heavily optimize our code if we use segment trees in finding out the minimum value up to the depth k above it. We can answer the query in logN time and reduce the overall complexity to 

O(Q + M*logN) which will pass easily.

How to use Segment Trees? 

We will apply dfs from the root node. Each time we iterate over all possible tickets, we can use query function to know the minimum possible value above k depths. We will build the segment tree on the depth. Once, we finish iterating a particular subtree, i.e that node and dfs over all it's children, we will update the depth to infinite so that it does not affect the answer for other branches of my tree.

Initially, we have not started iteration on any nodes, so the answer of dp[x] will be infinite and tree[ind] will also be infinite.  So, we don't need to build a tree here, because all values are initially infinite.

Now, dp[1] will be zero. Because we don't need any cost to visit node 1 if we are at node 1 itself.

So, we will also update the tree with dp[1] value. Note that the depth of node 1 is 0. 

The code:

ll N,m; vector<ll> adj[100005]; vector<pll> tickets[100005]; ll tree[400025]; ll dp[100005]; void update(ll ind, ll l, ll r, ll t, ll val){ if (t < l || t > r){ return; } if (l == r){ tree[ind] = val; return; } ll mid = (l+r)/2; update(2*ind,l,mid,t,val); update(2*ind+1,mid+1,r,t,val); tree[ind] = min(tree[2*ind],tree[2*ind+1]); } ll query(ll ind, ll l, ll r, ll qs, ll qe){ if (qe < l || qs > r){ return INF; } if (l >= qs && r <= qe){ return tree[ind]; } ll mid = (l+r)/2; return min(query(2*ind,l,mid,qs,qe),query(2*ind+1,mid+1,r,qs,qe)); } void dfs(ll node, ll depth){ for(auto i:tickets[node]){ ll k = i.F; ll w = i.S; dp[node] = min(dp[node],query(1,0,N,max((1ll)*0,depth-k),depth)+w); } update(1,0,N,depth,dp[node]); for(ll child:adj[node]){ dfs(child,depth+1); } update(1,0,N,depth,INF); } void solve() { cin >> N >> m; rep(i,0,N-1){ ll a,b; cin >> a >> b; adj[b].PB(a); } rep(i,0,m){ ll v,d,w; cin >> v >> d >> w; tickets[v].PB({d,w}); } rep(i,0,400025) tree[i] = INF; rep(i,0,100005) dp[i] = INF; dp[1] = 0; update(1,0,N,0,0); dfs(1,0); ll q; cin >> q; rep(i,0,q){ ll node; cin >> node; cout << dp[node] << endl; } } int32_t main(){ ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0); int t=1; // cin>>t; test(t){ solve(); } return 0; }