Given an n-ary tree with weights consisting of N nodes.
We are given Q queries. each query has 3 integers of following form:-
node1
node2
val
For each such query, we want to find out number of nodes which belong to path between node1
and node2
and have weights less than val
?
Constraints
(Will be updating other questions and details)
Let F(U, K)
return number of edges from node U
to root node
which have weight less than K
.
Then,
Required ans = number of edges with weight less than K that fall in path from node A to B
= F(A, K) + F(B, K) - F(LCA(A,B), K)
Here LCA(A, B)
is Lowest Common Ancestor of A
, B
nodes.
We have subtracted it to remove double counting in case where LCA
of A
, B
isn't root (illustrated in picture by Case #2)
Solution #1
This is straightforward to do with HLD (Heavy Light Decomposition), which is a datastructure used heavily in CP. One can copy and paste HLD code template to solve it quickly.
But this isn't required, there exists another simpler solution.
... drumrolls 🥁 ...
Solution #2
Firstly process inputs by doing :-
Now imagine the tree with each edge weight = 0. We will keep solving each query in the sorted order. To solve query-1 with K_1, we process all edges with edge weight less than K_1.
How to process an edge with weight W?
We update edge weight in our new tree which had initially 0 weights. How will this affect our tree? We need to update the subtree below current edge, as all their F(U, Z) with Z >= W will get updated.
How do we add something to a subtree?
This was taught in Trees class. You can do so by using Euler Tour concept and using any Range Query you like (Segment Tree, Fenwick Tree, etc).
This solves our problem in O(M + Q + M x log N)
.