Reputation: 1316
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
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
Upvotes: 1
Views: 143
Reputation: 19
I found example on querying fenwick tree's in princeton algorithms. You can find it here:
http://algs4.cs.princeton.edu/99misc/FenwickTree.java.html
Hope it will help you for your context.
Upvotes: 1
Reputation: 3818
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.
Upvotes: 1