Mayur Kulkarni
Mayur Kulkarni

Reputation: 1316

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

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

Answers (2)

Nadja
Nadja

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

delta
delta

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

Related Questions