#P3848. 第2题-挂满彩灯的二叉树
-
1000ms
Tried: 60
Accepted: 27
Difficulty: 5
所属公司 :
拼多多
时间 :2025年9月28日
-
算法标签>BFS
第2题-挂满彩灯的二叉树
解题思路
设彩灯颜色在集合 {1,2,3,4,5} 上循环。按某个节点一次,会使该节点及其整棵子树的颜色都加 1(模 5)。
把颜色映射为 0…4 更便于计算:令 c0 = c - 1,则每次点击相当于对子树的颜色 “+1(mod 5)”。
若从根到当前节点一共累计了 inc 次加 1(mod 5),则当前节点颜色为 (initial0 + inc) % 5。要让它变成 target0,最少还需在当前节点再按
d = (target0 - initial0 - inc) mod 5 (取 0..4 之间)
次。这样既把当前节点修正为目标色,又把「增量」传递给整棵子树。之后访问它的孩子时,将路径累计增量更新为 inc' = (inc + d) % 5。
关键性质(贪心正确性):
- 祖先的点击会影响当前节点,但后代的点击不会影响当前节点。因此当走到某节点时,祖先对它的影响已经定型,想要把该节点修正到目标色,唯一能动用的只有它自己(及其祖先,祖先已过时机)。因此“自顶向下、每点就地修正为目标”的贪心是全局最优。
- 每个节点的最少点击次数就是上述
d∈{0,1,2,3,4},总答案为所有访问到的非空节点的d之和。
实现上用迭代 DFS/BFS(避免递归爆栈),在数组(层序存储)的二叉树上自顶向下遍历。遇到 initial[i]==0 视为该位置无节点,直接跳过,也不扩展其子节点。
算法要点:
- 将每个颜色减 1,转为 0…4。
- 维护栈/队列元素
(idx, inc):节点下标与从根到此的累计加值。 - 计算
d,累加到答案,并把(child_idx, (inc + d) % 5)加入待处理。 - 超界或空节点(0)直接跳过。
复杂度分析
- 时间复杂度:O(n),每个非空节点至多入栈出栈一次,常数操作。
- 空间复杂度:O(h) 平均(树高),最坏 O(n)(用显式栈/队列;对完全二叉树也在可接受范围)。 不使用递归,避免深树时的栈溢出。
代码实现
Python
# 题面功能写在外部函数里
def min_clicks(initial, target):
n = len(initial)
ans = 0
# 映射到 0..4,方便取模
a = [x - 1 if x != 0 else 0 for x in initial]
b = [x - 1 if x != 0 else 0 for x in target]
# 迭代 DFS 栈:元素为 (index, inc)
stack = [(0, 0)]
while stack:
i, inc = stack.pop()
if i >= n or initial[i] == 0:
continue
# 当前节点需要的点击次数 d(取 0..4)
d = (b[i] - a[i] - inc) % 5
ans += d
inc2 = (inc + d) % 5
# 左右孩子
l = 2 * i + 1
r = 2 * i + 2
if l < n:
stack.append((l, inc2))
if r < n:
stack.append((r, inc2))
return ans
# 主函数:ACM 风格输入输出
def main():
import sys
data = sys.stdin.read().strip().split()
it = iter(data)
n = int(next(it))
initial = [int(next(it)) for _ in range(n)]
target = [int(next(it)) for _ in range(n)]
print(min_clicks(initial, target))
if __name__ == "__main__":
main()
Java
import java.io.*;
import java.util.*;
// 类名必须为 Main(ACM 风格)
public class Main {
// 外部函数:计算最少点击次数
static long minClicks(int[] initial, int[] target) {
int n = initial.length;
long ans = 0L;
// 转为 0..4
int[] a = new int[n];
int[] b = new int[n];
for (int i = 0; i < n; i++) {
a[i] = initial[i] == 0 ? 0 : initial[i] - 1;
b[i] = target[i] == 0 ? 0 : target[i] - 1;
}
// 迭代 DFS,自定义整型栈(避免装箱开销,n 可达 1e6)
int[] idxStack = new int[Math.max(1, n)];
int[] incStack = new int[Math.max(1, n)];
int top = 0;
idxStack[top] = 0;
incStack[top] = 0;
top++;
while (top > 0) {
top--;
int i = idxStack[top];
int inc = incStack[top];
if (i >= n || initial[i] == 0) continue;
int d = (b[i] - a[i] - inc) % 5;
if (d < 0) d += 5; // Java 取模可能为负,修正到 0..4
ans += d;
int inc2 = (inc + d) % 5;
int l = 2 * i + 1, r = 2 * i + 2;
if (l < n) {
idxStack[top] = l;
incStack[top] = inc2;
top++;
}
if (r < n) {
idxStack[top] = r;
incStack[top] = inc2;
top++;
}
}
return ans;
}
// 快速读入:数据规模大(n 可达 1e6),用字节流更稳妥
static class FastScanner {
private final InputStream in;
private final byte[] buffer = new byte[1 << 16];
private int ptr = 0, len = 0;
FastScanner(InputStream is) { in = is; }
private int read() throws IOException {
if (ptr >= len) {
len = in.read(buffer);
ptr = 0;
if (len <= 0) return -1;
}
return buffer[ptr++];
}
int nextInt() throws IOException {
int c, sgn = 1, x = 0;
do { c = read(); } while (c <= 32);
if (c == '-') { sgn = -1; c = read(); }
while (c > 32) {
x = x * 10 + (c - '0');
c = read();
}
return x * sgn;
}
}
// 主函数:ACM 风格
public static void main(String[] args) throws Exception {
FastScanner fs = new FastScanner(System.in);
int n = fs.nextInt();
int[] initial = new int[n];
int[] target = new int[n];
for (int i = 0; i < n; i++) initial[i] = fs.nextInt();
for (int i = 0; i < n; i++) target[i] = fs.nextInt();
System.out.println(minClicks(initial, target));
}
}
C++
#include <bits/stdc++.h>
using namespace std;
// 外部函数:计算最少点击次数
long long min_clicks(const vector<int>& initial, const vector<int>& target) {
int n = (int)initial.size();
long long ans = 0;
// 转为 0..4
vector<int> a(n), b(n);
for (int i = 0; i < n; ++i) {
a[i] = (initial[i] == 0) ? 0 : (initial[i] - 1);
b[i] = (target[i] == 0) ? 0 : (target[i] - 1);
}
// 迭代 DFS 栈:(idx, inc)
vector<int> idx, inc;
idx.reserve(n);
inc.reserve(n);
idx.push_back(0);
inc.push_back(0);
while (!idx.empty()) {
int i = idx.back(); idx.pop_back();
int cur = inc.back(); inc.pop_back();
if (i >= n || initial[i] == 0) continue;
int d = (b[i] - a[i] - cur) % 5;
if (d < 0) d += 5; // C++ 负模修正
ans += d;
int cur2 = (cur + d) % 5;
int l = 2 * i + 1;
int r = 2 * i + 2;
if (l < n) { idx.push_back(l); inc.push_back(cur2); }
if (r < n) { idx.push_back(r); inc.push_back(cur2); }
}
return ans;
}
// 主函数:ACM 风格输入输出
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n;
if (!(cin >> n)) return 0;
vector<int> initial(n), target(n);
for (int i = 0; i < n; ++i) cin >> initial[i];
for (int i = 0; i < n; ++i) cin >> target[i];
cout << min_clicks(initial, target) << "\n";
return 0;
}
题目内容
圣诞节快到了,有一棵挂满彩灯的二叉树,需要你来按照图纸装饰。彩灯有 5 种颜色变化,分别用 1−5 表示。
1 表示红色,2 表示黄色,3 表示蓝色,4 表示紫色,5 表示绿色。
每个节点都一个颜色控制器,每按一下,就会将当前彩灯以及以当前节点为根节点的子树上的所有节点,切换到下一个颜色(红->黄 ->蓝->紫>绿->红...)循环切换。
给定二叉树的初始状态 initial 和 目标状态 target ,两者都以层序遍历产出的一维数组表示。数组元素对应对应位置节点的颜色,0 表示该节点没有彩灯。
请给出从 initial 状态切换至 target 状态需要的最少控制器点击次数。
注意: 控制器按一下,不只影响当前节点,也会影响以当前节点为根节点的子树上所有节点切换到下一个颜色(最终不一定是同一个颜色)。
输入描述
第一行输入为一个整数 n ,代表 intial 和 target 数组的大小。
第二行输入为 n 个整数,代表 initial 数组。
第三行输入为 n 个整数,如果 initial[i]==0 , 则 targt[i] 也一定为 0 。
1<=initial.length<=106
输出描述
一个整数,表示最少点击次数
样例1
输入
5
1 2 3 0 1
2 3 1 0 2
输出
3
样例2
输入
7
1 2 3 1 2 3 1
3 1 2 3 1 2 1
输出
10