java - Fenwick 树中的范围查询

Question

我正在尝试解决这个涉及查询分域树的问题。问题陈述如下:

这是一个来自 Hackerrank 竞赛的问题:link

给定一棵树，其中每个节点都标记为 1、2、...、n。这棵树中有多少相似对(S)？

一对(A,B)是相似对当且仅当

• 节点A是节点B的祖先
• abs(A - B) <= T.

输入格式:输入的第一行包含两个整数 n 和 T。接下来是 n-1 行，每行包含两个整数 si 和 ei，其中节点 si 是节点 ei 的父节点。

输出格式:输出一个整数，表示树中相似对的数量

约束:

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

也保证没有循环，但树不一定是二叉树。

示例输入:

5 2
3 2
3 1
1 4
1 5

示例输出:

4

解释:相似对是:(3, 2) (3, 1) (3, 4) (3, 5)

我的算法:我会在DFS中遍历树，同时维护一个HashSet S用于查询。在进入节点时，我会将节点 x 添加到 S，而在离开节点时，我会将其从集合中移除。

现在为了找到特定叶节点的答案，我需要找出 Set 中遵循 x-T <= y <= x+T 的节点数，其中“y”是祖先，x-T 和 x+T 是范围内的节点。

我理解 Fenwick 树的概念，但我无法设计在树中存储什么以便进行范围查询，或者具体如何进行范围查询。我了解检索总和时的工作原理，但是给定一个范围 [left, right] 我如何找到存储在树中的范围内的元素数

Answer 1

注意label的范围最多为10^5，所以可以将label存储在Fenwick tree中。用1表示存在节点，否则为0。则sum表示节点总数。

假设我们有 Fenwick 树 T，有方法

• add(pos, val), add val 位置 pos, O(lgn)<
• sum(pos)，获取位置1到pos的和，O(lgn)

而dfs，当插入节点x时，做T.add(x, 1)，当删除节点x时，做T.add(x, -1)，对节点x进行查询时，为T.sum(x+k) - T.sum(x-k -1).

参见 code了解更多详情。