【备战秋招】每日一题:第六题-二叉树染色(Ⅱ)
在线评测链接:P1105
题目内容
塔子哥是一个热爱编程的大学生，他喜欢参加各种算法竞赛，挑战自己的智力。有一天，他在网上看到了一个有趣的二叉树问题，他决定尝试一下。他用纸和笔画出了一个 n n n 层的满二叉树（共 2 n − 1 2^n-1 2n−1 个节点，编号从 1 1 1 到 2 n − 1 2^n-1 2n−1 ，对于编号为 i i i （ 1 ≤ i ≤ 2 n − 1 − 1 1 \le i\le 2^{n-1} -1 1≤i≤2n−1−1 ）的节点，它的左儿子为 2 ∗ i 2*i 2∗i ，它的右儿子为 2 ∗ i + 1 2*i+1 2∗i+1），然后用红色的笔操作 q q q 次，每次都是在某些节点上画圈，表示将这些节点及其子树染红。他想知道每次染红后，二叉树中有多少个红色的节点。
他觉得这个问题很简单，只要用递归或者栈就可以解决。但是，当他开始编写代码时，他发现问题并没有那么容易。他需要考虑很多细节，比如如何存储和更新二叉树的状态，如何避免重复计算，如何优化时间和空间复杂度等等。他越写越觉得头疼，越写越觉得不对劲。他开始怀疑自己的能力，甚至怀疑这个问题是否有解。他不甘心放弃，他决定向你求助，希望你能给他一些提示或者思路。你能帮助塔子哥吗？
输入描述
第一行输入两个正整数 n n n 和 q q q ，代表二叉树的层数和操作次数。
接下来的 q q q 行，每行输入一个正整数 a i a_i ai ，代表染色的节点编号。
1 ≤ n ≤ 40 1\le n \le 40 1≤n≤40
1 ≤ q ≤ 10000 1 \le q \le 10000 1≤q≤10000
1 ≤ a i ≤ 2 n 1 \le a_i \le 2^n 1≤ai≤2n
输出描述
输出 q q q 行，每行输入一个正整数，代表当前操作结束后二叉树的红色节点数量。
样例
输入
4 2
4
3
输出
3
10
思路
巧妙优化
强烈建议对照着代码理解下方的思路!
step1:观察到这棵树节点非常多 , 达到 2 40 2^{40} 240 这个量级。所以不可能真的去建一颗树去模拟这个过程，暴力不可行。所以需要考虑优化：
假设上一轮操作之后的红色节点个数为 r e s res res ， 查询的节点序列为 { x 1 , x 2 , . . . , x q } \{x_1,x_2,...,x_q\} {x1,x2,...,xq}
1.如果第 i i i轮的查询节点 x i x_i xi已经被染色过了，那么节点个数不变，直接输出 r e s res res。 这个只需要判断 { x 1 , x 2 , . . . , x i − 1 } \{x_1,x_2,...,x_{i-1}\} {x1,x2,...,xi−1} 这个集合是否存在 x i x_i xi的祖先。
所以不难想到的是，用一个哈希表 H H H来装当前已经查询的节点集合，然后不断的除2： x i : = x i 2 x_i := \frac{x_i}{2} xi:=2xi 来模拟从 x i x_i xi 跳到根节点的过程 ，查询 H H H是否存在该值即可。
这个部分的复杂度为 O ( d e p t h ( x i ) ) = O ( n ) O(depth(x_i))=O(n) O(depth(xi))=O(n)
2.如果第 i i i轮的查询节点 x i x_i xi未被染色过，那么就用满二叉子树的个数 减去 子树中红色节点个数。
暴力做法:一个容易想到的做法是既然我们有了 H H H了。那么枚举 H H H 中的节点，判断它们和 x i x_i xi 的关系。只要是子树关系就减去。 但这样的单次询问复杂度是 O ( n q ) O(nq) O(nq),总复杂度是 O ( n q 2 ) O(nq^2) O(nq2)的. 虽然数据随机的情况下，期望复杂度可能没这么大，但是最好还是寻找更优的做法！
更优做法：上述的暴力做法实际上是存在重复计算的 ， 根据这一点，我们可以想到，不用等到必须要求子树中的红色节点的个数的时候再去求。而是提前就更新好。
具体的，我们维护一个哈希表 c n t cnt cnt，来存当前每个点的子树的红色节点个数。当一次查询结束时，我们用这个点来更新所有祖先节点,
c n t [ f a ] : = c n t [ f a ] + r e s cnt[fa] := cnt[fa]+res cnt[fa]:=cnt[fa]+res
f a fa fa是当前节点的所有祖先节点,这里也是类似第一步地反复跳父亲节点。
这样我们保证了在查询第 i i i个节点的时候，每个之前的节点 x j , j ∈ [ 1 , i − 1 ] x_j,j \in [1,i-1] xj,j∈[1,i−1]的贡献都被计算在树上了。那么 c n t [ x i ] cnt[x_i] cnt[xi] 直接就是以 x i x_i xi 为节点的红色节点个数。那么第 i i i次的贡献就是
Δ i = 2 n − d e p t h ( x i ) + 1 − 1 + c n t [ x i ] \Delta_i=2^{n-depth(x_i)+1}-1+cnt[x_i] Δi=2n−depth(xi)+1−1+cnt[xi]
, 答案就是 c n t [ r o o t ] = c n t [ 1 ] cnt[root]=cnt[1] cnt[root]=cnt[1] . 而计算 Δ i \Delta_i Δi 的开销就是 O ( 1 ) O(1) O(1)的了！
所以总时间复杂度为: O ( n q ) O(nq) O(nq)
这里注意，对于 c n t cnt cnt , 由于一个点最多有 n n n个祖先。所以 q q q次查询最多涉及到 n q nq nq个不同的结点。所以哈希表 c n t cnt cnt 的空间复杂度是 O ( n q ) O(nq) O(nq) 的，并不会导致空间爆炸
类似题目推荐
本题有一定的难度！这里推荐几道类似的不能直接模拟，需要一定优化技巧的题
CodeFun2000
P1179 2023.04.09-小红书-春招-第三题-神奇的盒子 - 十分相似！
P1122 2023.03.26-小红书-第三题-涂色 - 思考不去套高级数据结构做本题，同样需要优化技巧
P1238 2023.04.15-美团实习-第四题-俄罗斯套娃 - 贪心 + multiset 优化
P1269 2023.04.29-美团春招-第四题-SSTF算法 - 不能直接模拟，需要multiset 优化
代码
CPP
#include python
from collections import defaultdict
n , q = list(map(int,input().split()))
s = set()
cnt = defaultdict(int)
for _ in range (q):
x = int(input())
# 如果已经被覆盖，就不用了
ok = True
y = x
while y != 0:
if y in s:
ok = False
y >>= 1
if not ok:
print (cnt[1])
continue
# 加入集合
s.add(x)
# 计算深度
dep = 0
y = x
while y != 0:
dep += 1
y >>= 1
# 满二叉的情况
fv = (1 << (n - dep + 1)) - 1
# 新增的个数
delta = fv - cnt[x]
# 更新祖先
y = x
while y != 0:
cnt[y] += delta
y >>= 1
print (cnt[1])
Java
import java.util.*;
public class Main {
public static void main(String[] args) {
Scanner sc = new Scanner(System.in);
int n = sc.nextInt(), q = sc.nextInt();
Set<Long> mp1 = new HashSet<>(); //用于存储之前涂过的所有点
Map<Long, Long> mp2 = new HashMap<>(); //用于存储每个点的子树中有多少点被涂完了
long sum = 0; //记录红色节点个数
for (int i = 0; i < q; i++) {
int deep = 0;
long t = sc.nextLong(); //当前打算涂的节点
for(long x = t; x != 0; x >>= 1){ //遍历当前节点的所有祖先节点
if(mp1.contains(x)) { //如果第一个哈希表中存在这个节点的祖先，当前节点就不需要处理，用deep作一下标记并退出循环
deep = -1;
break;
}
deep++;
}
if(deep == -1) { //不需要处理当前节点
System.out.println(sum);
continue;
}
mp1.add(t); //将当前节点存入第一个哈希表中
long num = (1L << n-deep + 1) - 1; //计算当前节点的子树节点个数
if(mp2.containsKey(t)) num -= mp2.get(t); //如果第二个哈希表中有当前节点，就减去之前已经涂过的节点个数
for(long x = t; x != 0; x >>= 1){ //将当前涂的数量更新到第二个哈希表中
if(mp2.containsKey(x))
mp2.put(x, num + mp2.get(x));
else
mp2.put(x, num);
}
sum += num;
System.out.println(sum);
}
}
}
Go
package main
import (
"fmt"
)
func main() {
var n, q int
fmt.Scan(&n, &q)
// 记录所有被染过色的节点
nodes := make(map[int]struct{})
// 每个节点的子树中被染色的节点个数
memo := make(map[int]int)
var ans int
for i := 0; i < q; i++ {
var root int
fmt.Scan(&root)
var deep int
for j := root; j > 0; j >>= 1 {
deep++
if _, ok := nodes[j]; ok {
// 有祖先节点已经被染过色了，跳过当前节点
deep = -1
break
}
}
if deep == -1 {
fmt.Println(ans)
continue
}
nodes[root] = struct{}{}
num := (1 << (n - deep + 1)) - 1
if _num, ok := memo[root]; ok {
// 当前节点有子节点已经染过色
num -= _num
}
for j := root; j > 0; j >>= 1 {
if _, ok := memo[j]; ok {
memo[j] += num
} else {
memo[j] = num
}
}
ans += num
fmt.Println(ans)
}
}
Js
const input = [];
process.stdin.resume();
process.stdin.setEncoding('utf-8');
process.stdin.on('data', (data) => {
input.push(...data.trim().split('\n'));
});
process.stdin.on('end', () => {
let [n, q] = input[0].split(' ').map(Number);
var s = new Set();
var cnt = new Map();
for (let i = 1; i <= q; i++) {
let x = Number(input[i]);
let ok = true;
let y = x;
// 如果已经被覆盖，就不用了
while (y !== 0) {
if (s.has(Number(y))) {
ok = false;
break;
}
y = Math.floor(y / 2);
}
if (!ok) {
console.log(cnt.get(1));
continue;
}
// 加入集合H
s.add(x);
let dep = 0;
y = x;
// 求深度
while (y !== 0) {
dep++;
y = Math.floor(y / 2);
}
// 满二叉子树个数
let fv = ((2 ** (n - dep + 1)) - 1);
// 子树红色节点个数
let delta = fv - (cnt.get(Number(x)) || 0);
y = x;
// 更新祖先节点
while (y !== 0) {
cnt.set(Number(y), (cnt.get(Number(y)) || 0) + delta);
y = Math.floor(y / 2);
}
console.log(cnt.get(1));
}
});