LeetCode 周赛上分之旅 #47 前后缀分解结合单调栈的贡献问题
⭐️ 本文已收录到 AndroidFamily，技术和职场问题，请关注公众号 [彭旭锐] 和 BaguTree Pro 知识星球提问。
学习数据结构与算法的关键在于掌握问题背后的算法思维框架，你的思考越抽象，它能覆盖的问题域就越广，理解难度也更复杂。在这个专栏里，小彭与你分享每场 LeetCode 周赛的解题报告，一起体会上分之旅。
本文是 LeetCode 上分之旅系列的第 47 篇文章，往期回顾请移步到文章末尾~
LeetCode 周赛 364
T1. 最大二进制奇数（Easy）
- 标签：贪心
T2. 美丽塔 I（Medium）
- 标签：枚举、前后缀分解、单调栈
T3. 美丽塔 II（Medium）
- 标签：枚举、前后缀分解、单调栈
T4. 统计树中的合法路径数目（Hard）
- 标签：DFS、质数
T1. 最大二进制奇数（Easy）
https://leetcode.cn/problems/maximum-odd-binary-number/description/
题解（模拟）
简单模拟题，先计算 $1$ 的个数，将其中一个 $1$ 置于最低位，其它 $1$ 置于最高位：
class Solution {
fun maximumOddBinaryNumber(s: String): String {
val cnt = s.count { it == '1' }
return StringBuilder().apply {
repeat(cnt - 1) {
append("1")
}
repeat(s.length - cnt) {
append("0")
}
append("1")
}.toString()
}
}
class Solution:
def maximumOddBinaryNumber(self, s: str) -> str:
n, cnt = len(s), s.count("1")
return "1" * (cnt - 1) + "0" * (n - cnt) + "1"
class Solution {
public:
string maximumOddBinaryNumber(string s) {
int n = s.length();
int cnt = 0;
for (auto& e : s) {
if (e == '1') cnt++;
}
string ret;
for (int i = 0; i < cnt - 1; i++) {
ret.push_back('1');
}
for (int i = 0; i < n - cnt; i++) {
ret.push_back('0');
}
ret.push_back('1');
return ret;
}
};
复杂度分析：
- 时间复杂度：$O(n)$ 线性遍历；
- 空间复杂度：$O(1)$ 不考虑结果字符串。
T2. 美丽塔 I（Medium）
https://leetcode.cn/problems/beautiful-towers-i/description/
同 T3. 美丽塔 I
T3. 美丽塔 II（Medium）
https://leetcode.cn/problems/beautiful-towers-ii/description/
问题分析
初步分析：
- 问题目标： 构造满足条件的方案，使得数组呈现山状数组，返回元素和；
- 方案条件： 从数组的最大值向左侧为递减，向右侧也为递减。
思考实现：
- 在 T2. 美丽塔 I（Medium） 中的数据量只有 $1000$，我们可以枚举以每个点作为山峰（数组最大值）的方案，从山顶依次向两侧递减，使得当前位置不高于前一个位置，整体的时间复杂度是 $O(n^2)$；
- 在 T3. 美丽塔 II（Medium） 中数据量有 $10^5$，我们需要思考低于平方时间复杂度的方法。
思考优化：
以示例 [6,5,3,9,2,7] 为例，我们观察以 $3$ 和 $9$ 作为山顶的两个方案：
以 3 作为山顶：
3 3 |3 3| 2 2
以 9 作为山顶
3 3 |3 9| 2 2
可以发现：以 $3$ 作为山顶的左侧与以 $9$ 为山顶的右侧在两个方案之间是可以复用的，至此发现解决方法：我们可以分别预处理出以每个节点作为山顶的前缀和后缀的和：
- $pre[i]$ 表示以 $maxHeights[i]$ 作为山顶时左侧段的前缀和；
- $suf[i]$ 表示以 $maxHeights[i]$ 作为山顶时右侧段的后缀和。
那么，最佳方案就是 $pre[i] + suf[i] - maxHeight[i]$ 的最大值。 现在，最后的问题是如何以均摊 $O(1)$ 的时间复杂度计算出每个元素前后缀的和？
思考递推关系：
继续以示例 [6,5,3,9,2,7] 为例：
- 以 $6$ 为山顶，前缀为 $[6]$
- 以 $5$ 为山顶，需要保证左侧元素不大于 $5$，因此找到 $6$ 并修改为 $5$，前缀为 $[5, 5]$
- 以 $3$ 为山顶，需要保证左侧元素不大于 $3$，因此找到两个 $5$ 并修改为 $3$，前缀为 $[3, 3, 3]$
- 以 $9$ 为山顶，需要保证左侧元素不大于 $9$，不需要修改，前缀为 $[3, 3, 3, 9]$
- 以 $2$ 为山顶，需要保证左侧元素不大于 $2$，修改后为 $[2, 2, 2, 2, 2]$
- 以 $7$ 为山顶，需要保证左侧元素不大于 $7$，不需要修改，前缀为 $[2, 2, 2, 2, 2, 7]$
提高抽象程度：
观察以上步骤，问题的关键在于修改操作：由于数组是递增的，因此修改的步骤就是在「寻找小于等于当前元素 $x$ 的上一个元素」，再将中间的元素削减为 $x$。「寻找上一个更小元素」，这是单调栈的典型场景。
题解一（枚举）
枚举以每个元素作为山顶的方案：
class Solution {
fun maximumSumOfHeights(maxHeights: List<Int>): Long {
val n = maxHeights.size
var ret = 0L
for (i in maxHeights.indices) {
var curSum = maxHeights[i].toLong()
var pre = maxHeights[i]
for (j in i - 1 downTo 0) {
pre = min(pre, maxHeights[j])
curSum += pre
}
pre = maxHeights[i]
for (j in i + 1 ..< n) {
pre = min(pre, maxHeights[j])
curSum += pre
}
ret = max(ret, curSum)
}
return ret
}
}
class Solution:
def maximumSumOfHeights(self, maxHeights: List[int]) -> int:
n, ret = len(maxHeights), 0
for i in range(n):
curSum = maxHeights[i]
pre = maxHeights[i]
for j in range(i + 1, n):
pre = min(pre, maxHeights[j])
curSum += pre
pre = maxHeights[i]
for j in range(i - 1, -1, -1):
pre = min(pre, maxHeights[j])
curSum += pre
ret = max(ret, curSum)
return ret
class Solution {
public:
long long maximumSumOfHeights(vector<int>& maxHeights) {
int n = maxHeights.size();
long long ret = 0;
for (int i = 0; i < n; i++) {
long long curSum = maxHeights[i];
int pre = maxHeights[i];
for (int j = i + 1; j < n; j++) {
pre = min(pre, maxHeights[j]);
curSum += pre;
}
pre = maxHeights[i];
for (int j = i - 1; j >= 0; j--) {
pre = min(pre, maxHeights[j]);
curSum += pre;
}
ret = max(ret, curSum);
}
return ret;
}
};
复杂度分析：
- 时间复杂度：$O(n^2)$ 每个方案的时间复杂度是 $O(n)$，一共有 $n$ 种方案；
- 空间复杂度：$O(1)$ 仅使用常量级别空间。
题解二（前后缀分解 + 单调栈）
使用单点栈维护前后缀数组，为了便于边界计算，我们构造长为 $n + 1$ 的数组。以示例 [6,5,3,9,2,7] 为例：
0, 5, 6, 10, 4, 5
13, 8, 6, 2, 1, 0
class Solution {
fun maximumSumOfHeights(maxHeights: List<Int>): Long {
val n = maxHeights.size
val suf = LongArray(n + 1)
val pre = LongArray(n + 1)
// 单调栈求前缀
val stack = java.util.ArrayDeque<Int>()
for (i in 0 until n) {
// 弹出栈顶
while (!stack.isEmpty() && maxHeights[stack.peek()] > maxHeights[i]) {
stack.pop()
}
val j = if (stack.isEmpty()) -1 else stack.peek()
pre[i + 1] = pre[j + 1] + 1L * (i - j) * maxHeights[i]
stack.push(i)
}
// 单调栈求后缀
stack.clear()
for (i in n - 1 downTo 0) {
// 弹出栈顶
while (!stack.isEmpty() && maxHeights[stack.peek()] > maxHeights[i]) {
stack.pop()
}
val j = if (stack.isEmpty()) n else stack.peek()
suf[i] = suf[j] + 1L * (j - i) * maxHeights[i]
stack.push(i)
}
// 合并
var ret = 0L
for (i in 0 until n) {
ret = max(ret, pre[i + 1] + suf[i] - maxHeights[i])
}
return ret
}
}
class Solution:
def maximumSumOfHeights(self, maxHeights: List[int]) -> int:
n = len(maxHeights)
suf = [0] * (n + 1)
pre = [0] * (n + 1)
stack = []
# 单调栈求前缀
for i in range(n):
# 弹出栈顶
while stack and maxHeights[stack[-1]] > maxHeights[i]:
stack.pop()
j = stack[-1] if stack else -1
pre[i + 1] = pre[j + 1] + (i - j) * maxHeights[i]
stack.append(i)
# 单调栈求后缀
stack = []
for i in range(n - 1, -1, -1):
# 弹出栈顶
while stack and maxHeights[stack[-1]] > maxHeights[i]:
stack.pop()
j = stack[-1] if stack else n
suf[i] = suf[j] + (j - i) * maxHeights[i]
stack.append(i)
# 合并
ret = 0
for i in range(n):
ret = max(ret, pre[i + 1] + suf[i] - maxHeights[i])
return ret
class Solution {
public:
long long maximumSumOfHeights(vector<int>& maxHeights) {
int n = maxHeights.size();
vector<long long> suf(n + 1, 0);
vector<long long> pre(n + 1, 0);
stack<int> st;
// 单调栈求前缀
for (int i = 0; i < n; i++) {
// 弹出栈顶
while (!st.empty() && maxHeights[st.top()] > maxHeights[i]) {
st.pop();
}
int j = st.empty() ? -1 : st.top();
pre[i + 1] = pre[j + 1] + 1LL * (i - j) * maxHeights[i];
st.push(i);
}
// 单调栈求后缀
while (!st.empty()) st.pop();
for (int i = n - 1; i >= 0; i--) {
// 弹出栈顶
while (!st.empty() && maxHeights[st.top()] > maxHeights[i]) {
st.pop();
}
int j = st.empty() ? n : st.top();
suf[i] = suf[j] + 1LL * (j - i) * maxHeights[i];
st.push(i);
}
// 合并
long long ret = 0;
for (int i = 0; i < n; i++) {
ret = max(ret, pre[i + 1] + suf[i] - maxHeights[i]);
}
return ret;
}
};
复杂度分析：
- 时间复杂度：$O(n)$ 在一侧的计算中，每个元素最多如何和出栈 $1$ 次；
- 空间复杂度：$O(n)$ 前后缀数组空间。
T4. 统计树中的合法路径数目（Hard）
https://leetcode.cn/problems/count-valid-paths-in-a-tree/description/
这道题似乎比 T3 还简单一些。
问题分析
初步分析：
- 问题目标： 寻找满足条件的方案数；
- 问题条件： 路径 $[a, b]$ 上质数的数目有且仅有 $1$；
- 问题要素： 路径和 - 表示路径上质数的数目。
思考实现：
- 子问题： 对于以根节点 x 的原问题，可以分为 3 种情况：
- 左子树可以构造的方案数
- 右子树可以构造的方案数
- 如果根节点为质数：「从根到子树节点的路径和为 $0$ 的数目」与「从根到其它子树节点的路径和为 $0$ 的数目」的乘积（乘法原理）
题解（DFS）
构造 DFS 函数，子树的 DFS 返回值为两个值：
- $cnt0$：到子树节点和为 $0$ 的路径数；
- $cnt1$：到子树节点和为 $1$ 的路径数；
返回结果时：
- 如果根节点为质数，那么只能与 $cnt0$ 个路径和为 $1$ 的路径；
- 如果根节点为非质数，那么 $cnt0$ 个路径可以组成和为 $0$ 的路径，同理 $cnt1$ 个路径可以组成和为 $1$ 的路径。
在子树的计算过程中还会构造结果：
由于题目说明 $[a, b]$ 与 $[b, a]$ 是相同路径，我们可以记录当前子树左侧已经计算过的 $cnt0$ 和 $cnt1$ 的累加和，再与当前子树的 $cnt0$ 与 $cnt1$ 做乘法：
$ret += cnt0 * cnt[1] + cnt1 * cnt[0]$
class Solution {
companion object {
val U = 100000
val primes = LinkedList<Int>()
val isPrime = BooleanArray(U + 1) { true }
init {
isPrime[1] = false
for (i in 2 .. U) {
if (isPrime[i]) primes.add(i)
for (e in primes) {
if (i * e > U) break
isPrime[i * e] = false
if (i % e == 0) break
}
}
}
}
fun countPaths(n: Int, edges: Array<IntArray>): Long {
val graph = Array(n + 1) { LinkedList<Int>() }
for ((from, to) in edges) {
graph[from].add(to)
graph[to].add(from)
}
var ret = 0L
// return 0 和 1 的数量
fun dfs(i: Int, pre: Int): IntArray {
// 终止条件
var cnt = IntArray(2)
if (isPrime[i]) {
cnt[1] = 1
} else {
cnt[0] = 1
}
// 递归
for (to in graph[i]) {
if (to == pre) continue // 返祖边
val (cnt0, cnt1) = dfs(to, i)
// 记录方案
ret += cnt0 * cnt[1] + cnt1 * cnt[0]
// 记录影响
if (isPrime[i]) {
cnt[1] += cnt0
} else {
cnt[0] += cnt0
cnt[1] += cnt1
}
}
return cnt
}
dfs(1, -1) // 随机选择根节点
return ret
}
}
U = 100000
primes = deque()
isPrime = [True] * (U + 1)
isPrime[1] = False
for i in range(2, U + 1):
if isPrime[i]: primes.append(i)
for e in primes:
if i * e > U: break
isPrime[i * e] = False
if i % e == 0: break
class Solution:
def countPaths(self, n, edges):
graph = defaultdict(list)
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
ret = 0
def dfs(i, pre):
nonlocal ret # 修改外部变量
cnt = [0, 0]
# 终止条件
if isPrime[i]:
cnt[1] = 1
else:
cnt[0] = 1
for to in graph[i]:
if to == pre: continue # 返祖边
cnt0, cnt1 = dfs(to, i)
# 记录方案
ret += cnt0 * cnt[1] + cnt1 * cnt[0]
# 记录影响
if isPrime[i]:
cnt[1] += cnt0
else:
cnt[0] += cnt0
cnt[1] += cnt1
return cnt
dfs(1, -1) # 随机选择根节点
return ret
const int U = 100000;
list<int> primes;
bool isPrime[U + 1];
bool inited = false;
void init() {
if (inited) return;
inited = true;
memset(isPrime, true, sizeof(isPrime));
isPrime[1] = false;
for (int i = 2; i <= U; ++i) {
if (isPrime[i]) primes.push_back(i);
for (auto e : primes) {
if (i * e > U) break;
isPrime[i * e] = false;
if (i % e == 0) break;
}
}
}
class Solution {
public:
long long countPaths(int n, vector<vector<int>>& edges) {
init();
vector<list<int>> graph(n + 1);
for (const auto& edge : edges) {
int from = edge[0];
int to = edge[1];
graph[from].push_back(to);
graph[to].push_back(from);
}
long long ret = 0;
// return 0 和 1 的数量
function<vector<int>(int, int)> dfs = [&](int i, int pre) -> vector<int> {
// 终止条件
vector<int> cnt(2, 0);
if (isPrime[i]) {
cnt[1] = 1;
} else {
cnt[0] = 1;
}
// 递归
for (auto to : graph[i]) {
if (to == pre) continue; // 返祖边
vector<int> subCnt = dfs(to, i);
int cnt0 = subCnt[0];
int cnt1 = subCnt[1];
// 记录方案
ret += cnt0 * cnt[1] + cnt1 * cnt[0];
// 记录影响
if (isPrime[i]) {
cnt[1] += cnt0;
} else {
cnt[0] += cnt0;
cnt[1] += cnt1;
}
}
return cnt;
};
dfs(1, -1); // 随机选择根节点
return ret;
}
};
复杂度分析：
- 时间复杂度：预处理时间为 $O(U)$，建图时间 和 DFS 时间为 $O(n)$；
- 空间复杂度：预处理空间为 $O(U)$，模拟空间为 $O(n)$。
枚举质数
枚举法：枚举 $[2, n]$ ，判断它是不是质数，整体时间复杂度是 $O(n\sqrt{n})$
// 暴力求质数
fun getPrimes(max: Int): IntArray {
val primes = LinkedList<Int>()
for (num in 2..max) {
if (isPrime(num)) primes.add(num)
}
return primes.toIntArray()
}
// 质数判断
fun isPrime(num: Int): Boolean {
var x = 2
while (x * x <= num) {
if (num % x == 0) return false
x++
}
return true
}
Eratosthenes 埃氏筛：如果 $x$ 是质数，那么 $x$ 的整数倍 $2x$、$3x$ 一定不是质数。我们设
isPrime[i]表示 $i$ 是否为质数。从小开始遍历，如果 $i$ 是质数，则同时将所有倍数标记为合数，整体时间复杂度是 $O(nlgn)$
为什么要从 $x^2$, $2x^2$ 开始标记，而不是 $2x$, $3x$ 开始标记，因为 $2x$, $3x$ 已经被小于 $x$ 的质数标记过。
// 埃氏筛求质数
val primes = LinkedList<Int>()
val isPrime = BooleanArray(U + 1) { true }
for (i in 2..U) {
// 检查是否为质数，这里不需要调用 isPrime() 函数判断是否质数，因为它没被小于它的数标记过，那么一定不是合数
if (!isPrime[i]) continue
primes.add(i)
// 标记
var x = i * i
while (x <= U) {
isPrime[x] = false
x += i
}
}
Euler 欧氏线性筛：尽管我们从 $x^2$ 开始标记来减少重复标记，但埃氏筛还是会重复标记合数。为了避免重复标记，标记 $x$ 与 “小于等于 $x$ 的最小质因子的质数” 的乘积为合数，保证每个合数只被标记最小的质因子标记，整体时间复杂度是 $O(n)$
// 线性筛求质数
val primes = LinkedList<Int>()
val isPrime = BooleanArray(U + 1) { true }
for (i in 2..U) {
// 检查是否为质数，这里不需要调用 isPrime() 函数判断是否质数，因为它没被小于它的数标记过，那么一定不是合数
if (isPrime[i]) {
primes.add(i)
}
// 标记
for (e in primes) {
if (i * e > U) break
isPrime[i * e] = false
if (i % e == 0) break
}
}
推荐阅读
LeetCode 上分之旅系列往期回顾：
⭐️ 永远相信美好的事情即将发生，欢迎加入小彭的 Android 交流社群~
