LeetCode 周赛上分之旅 #47 前后缀分解结合单调栈的贡献问题

彭旭锐2023-09-24 20:32

⭐️ 本文已收录到 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 置于最高位：

Kotlin 复制代码 
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()
    }
}
Python 复制代码 
class Solution:
    def maximumOddBinaryNumber(self, s: str) -> str:
        n, cnt = len(s), s.count("1")
        return "1" * (cnt - 1) + "0" * (n - cnt) + "1"
C++ 复制代码 
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 作为山顶的两个方案：

bash 复制代码 
以 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。「寻找上一个更小元素」，这是单调栈的典型场景。

题解一（枚举）

枚举以每个元素作为山顶的方案：

Kotlin 复制代码 
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
    }
}
Python 复制代码 
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
C++ 复制代码 
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] 为例：

bash 复制代码 
0, 5, 6, 10, 4, 5
13, 8, 6, 2, 1, 0
Kotlin 复制代码 
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
    }
}
Python 复制代码 
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
C++ 复制代码 
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）

bash 复制代码 
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\]

Kotlin 复制代码 
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
    }
}
Python 复制代码 
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
C++ 复制代码 
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)。

枚举质数

OI - 素数筛法

枚举法：枚举 \[2, n\] ，判断它是不是质数，整体时间复杂度是 O(n\\sqrt{n})

kotlin 复制代码 
// 暴力求质数
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 的质数标记过。

kotlin 复制代码 
// 埃氏筛求质数
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)

kotlin 复制代码 
// 线性筛求质数
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 上分之旅系列往期回顾：

热门推荐
01新手学习VR全景需要知道的几个问题02Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO03玄机平台应急响应—webshell查杀04RAG 实践- Ollama+RagFlow 部署本地知识库05优化手机性能，解决卡顿问题：关闭这3个微信开关，释放内存空间06HCIA-datacom数通题库和录播视频资料07docker安装启动问题解决排查08【一文读懂】NTN（非地面网络）技术介绍09支付宝短视频平台创作分成激励项目10组基轨迹建模 GBTM的介绍与实现（Stata 或 R）