Mayur Kulkarni Mayur Kulkarni - 3 months ago 20
Java Question

Range querying in Fenwick Tree

I'm trying to solve this problem that involves querying a fenwick tree.
Problem statement is as follows:

This is a problem from a Hackerrank contest: link

You are given a tree where each node is labeled from 1, 2, …, n. How many similar pairs(S) are there in this tree?

A pair (A,B) is a similar pair iff


  • node A is the ancestor of node B

  • abs(A - B) <= T.



Input format:
The first line of the input contains two integers n and T. This is followed by n-1 lines each containing two integers si and ei where node si is a parent to node ei.

Output format:
Output a single integer which denotes the number of similar pairs in the tree

Constraints:

1 <= n <= 100000
0 <= T <= n
1 <= si, ei <= n.


It is also guaranteed there are no cycles, but the tree does not have to be a binary tree.

Sample Input:

5 2
3 2
3 1
1 4
1 5


Sample Output:

4


Explanation:
The similar pairs are: (3, 2) (3, 1) (3, 4) (3, 5)

My Algorithm:
I will traverse the tree in DFS while maintaining a HashSet S for querying. While entering the node I'll add a node x to S and while leaving I'll remove it from set.

Now inorder to find the answer at a particular leaf node I need to find out the number of nodes in Set that follow x-T <= y <= x+T where 'y' is ancestor and x-T and x+T are the nodes in the range.

I understand the concept of Fenwick Tree but I'm unable to devise what to store in the tree inorder to make the ranged query, or how to make the ranged query in particular. I understand how it works when retrieving the sum, but given a range [left, right] how do I find the number of elements in the range stored in tree

Answer

Note the range of label is at most 10^5, so you can store the label in Fenwick tree. Using 1 to indicate the existence of a node, otherwise 0. Then sum indicates the total number of nodes.

Assume we have Fenwick tree T, with methods

  • add(pos, val), add val in position pos, O(lgn)
  • sum(pos), get the sum of position 1 to pos, O(lgn)

While dfs, when inserting a node x, do T.add(x, 1), when removing a node x, do T.add(x, -1), when do query on node x, it is T.sum(x+k) - T.sum(x-k-1).

See the code for more details.