LeetCode 双周赛 103（2023/04/29）区间求和的树状数组经典应用

本文已收录到 AndroidFamily，技术和职场问题，请关注公众号 [彭旭锐] 提问。

大家好，我是小彭。

这场周赛是 LeetCode 双周赛第 103 场，难得在五一假期第一天打周赛的人数也没有少太多。这场比赛前 3 题比较简单，我们把篇幅留给最后一题。

往期周赛回顾：LeetCode 单周赛第 342 场 · 容斥原理、计数排序、滑动窗口、子数组 GCB

周赛概览

Q1. K 个元素的最大和（Easy）

简单模拟题，不过多讲解。

Q2. 找到两个数组的前缀公共数组（Medium）

简单模拟题，在计数的实现上有三种解法：

• 解法 1：散列表 $O(n)$ 空间复杂度
• 解法 2：技数数组 $O(n)$ 空间复杂度
• 解法 3：状态压缩 $O(1)$ 空间复杂度

Q3. 网格图中鱼的最大数目（Hard）

这道题的难度标签是认真的吗？打 Medium 都过分了居然打 Hard？

• 解法 1：BFS / DFS $O(nm)$
• 解法 2：并查集 $O(nm)$

Q4. 将数组清空（Hard）

这道题的难点在于如何想到以及正确地将原问题转换为区间求和问题，思路想清楚后用树状数组实现。

• 解法 1：树状数组 + 索引数组 $O(nlgn)$
• 解法 2：树状数组 + 最小堆 $O(nlgn)$

---

Q1. K 个元素的最大和（Easy）

https://leetcode.cn/problems/maximum-sum-with-exactly-k-elements/

题目描述

给你一个下标从 0 开始的整数数组 nums 和一个整数 k 。你需要执行以下操作 恰好 k 次，最大化你的得分：

1. 从 nums 中选择一个元素 m 。
2. 将选中的元素 m 从数组中删除。
3. 将新元素 m + 1 添加到数组中。
4. 你的得分增加 m 。

请你返回执行以上操作恰好 k 次后的最大得分。

示例 1：

输入：nums = [1,2,3,4,5], k = 3
输出：18
解释：我们需要从 nums 中恰好选择 3 个元素并最大化得分。
第一次选择 5 。和为 5 ，nums = [1,2,3,4,6] 。
第二次选择 6 。和为 6 ，nums = [1,2,3,4,7] 。
第三次选择 7 。和为 5 + 6 + 7 = 18 ，nums = [1,2,3,4,8] 。
所以我们返回 18 。
18 是可以得到的最大答案。

示例 2：

输入：nums = [5,5,5], k = 2
输出：11
解释：我们需要从 nums 中恰好选择 2 个元素并最大化得分。
第一次选择 5 。和为 5 ，nums = [5,5,6] 。
第二次选择 6 。和为 6 ，nums = [5,5,7] 。
所以我们返回 11 。
11 是可以得到的最大答案。

提示：

• 1 <= nums.length <= 100
• 1 <= nums[i] <= 100
• 1 <= k <= 100

预备知识 - 等差数列求和

• 等差数列求和公式：(首项 + 尾项) * 项数 / 2

题解（模拟 + 贪心）

显然第一次操作的分数会选择数组中的最大值 max，后续操作是以 max 为首项的等差数列，直接使用等差数列求和公式即可。

class Solution {
    fun maximizeSum(nums: IntArray, k: Int): Int {
        val max = Arrays.stream(nums).max().getAsInt()
        return (max + max + k - 1) * k / 2
    }
}

复杂度分析：

• 时间复杂度：$O(n)$ 其中 n 是 nums 数组的长度；
• 空间复杂度：$O(1)$

---

Q2. 找到两个数组的前缀公共数组（Medium）

https://leetcode.cn/problems/find-the-prefix-common-array-of-two-arrays/

题目描述

给你两个下标从 0 开始长度为 n 的整数排列 A 和 B 。

A 和 B 的 前缀公共数组 定义为数组 C ，其中 C[i] 是数组 A 和 B 到下标为 i 之前公共元素的数目。

请你返回 A 和 B 的 前缀公共数组 。

如果一个长度为 n 的数组包含 1 到 n 的元素恰好一次，我们称这个数组是一个长度为 n 的 排列 。

示例 1：

输入：A = [1,3,2,4], B = [3,1,2,4]
输出：[0,2,3,4]
解释：i = 0：没有公共元素，所以 C[0] = 0 。
i = 1：1 和 3 是两个数组的前缀公共元素，所以 C[1] = 2 。
i = 2：1，2 和 3 是两个数组的前缀公共元素，所以 C[2] = 3 。
i = 3：1，2，3 和 4 是两个数组的前缀公共元素，所以 C[3] = 4 。

示例 2：

输入：A = [2,3,1], B = [3,1,2]
输出：[0,1,3]
解释：i = 0：没有公共元素，所以 C[0] = 0 。
i = 1：只有 3 是公共元素，所以 C[1] = 1 。
i = 2：1，2 和 3 是两个数组的前缀公共元素，所以 C[2] = 3 。

提示：

• 1 <= A.length == B.length == n <= 50
• 1 <= A[i], B[i] <= n
• 题目保证 A 和 B 两个数组都是 n 个元素的排列。

题解一（散列表）

从左到右遍历数组，并使用散列表记录访问过的元素，以及两个数组交集：

class Solution {
    fun findThePrefixCommonArray(A: IntArray, B: IntArray): IntArray {
        val n = A.size
        val ret = IntArray(n)
        val setA = HashSet<Int>()
        val setB = HashSet<Int>()
        val interSet = HashSet<Int>()
        for (i in 0 until n) {
            setA.add(A[i])
            setB.add(B[i])
            if (setB.contains(A[i])) interSet.add(A[i])
            if (setA.contains(B[i])) interSet.add(B[i])
            ret[i] = interSet.size
        }
        return ret
    }
}

复杂度分析：

• 时间复杂度：$O(n)$ 其中 n 是 nums 数组的长度；
• 空间复杂度：$O(n)$ 散列表空间。

题解二（计数数组）

题解一需要使用多倍空间，我们发现 A 和 B 都是 n 的排列，当访问到的元素 nums[i] 出现 2 次时就必然处于数组交集中。因此，我们不需要使用散列表记录访问过的元素，而只需要记录每个元素出现的次数。

class Solution {
    fun findThePrefixCommonArray(A: IntArray, B: IntArray): IntArray {
        val n = A.size
        val ret = IntArray(n)
        val cnt = IntArray(n + 1)
        var size = 0
        for (i in 0 until n) {
            if (++cnt[A[i]] == 2) size ++
            if (++cnt[B[i]] == 2) size ++
            ret[i] = size
        }
        return ret
    }
}

复杂度分析：

• 时间复杂度：$O(n)$ 其中 n 是 nums 数组的长度；
• 空间复杂度：$O(n)$ 计数数组空间；

题解三（状态压缩）

既然 A 和 B 的元素值不超过 50，我们可以使用两个 Long 变量代替散列表优化空间复杂度。

class Solution {
    fun findThePrefixCommonArray(A: IntArray, B: IntArray): IntArray {
        val n = A.size
        val ret = IntArray(n)
        var flagA = 0L
        var flagB = 0L
        var size = 0
        for (i in 0 until n) {
            flagA = flagA or (1L shl A[i])
            flagB = flagB or (1L shl B[i])
            // Kotlin 1.5 才有 Long.countOneBits()
            // ret[i] = (flagA and flagB).countOneBits()
            ret[i] = java.lang.Long.bitCount(flagA and flagB)
        }
        return ret
    }
}

复杂度分析：

• 时间复杂度：$O(n)$ 其中 n 是 nums 数组的长度；
• 空间复杂度：$O(1)$ 仅使用常量级别空间；

---

Q3. 网格图中鱼的最大数目（Hard）

https://leetcode.cn/problems/maximum-number-of-fish-in-a-grid/description/

题目描述

给你一个下标从 0 开始大小为 m x n 的二维整数数组 grid ，其中下标在 (r, c) 处的整数表示：

• 如果 grid[r][c] = 0 ，那么它是一块 陆地 。
• 如果 grid[r][c] > 0 ，那么它是一块 水域 ，且包含 grid[r][c] 条鱼。

一位渔夫可以从任意 水域 格子 (r, c) 出发，然后执行以下操作任意次：

• 捕捞格子 (r, c) 处所有的鱼，或者
• 移动到相邻的 水域 格子。

请你返回渔夫最优策略下， 最多 可以捕捞多少条鱼。如果没有水域格子，请你返回 0 。

格子 (r, c) 相邻 的格子为 (r, c + 1) ，(r, c - 1) ，(r + 1, c) 和 (r - 1, c) ，前提是相邻格子在网格图内。

示例 1：

输入：grid = [[0,2,1,0],[4,0,0,3],[1,0,0,4],[0,3,2,0]]
输出：7
解释：渔夫可以从格子(1,3) 出发，捕捞 3 条鱼，然后移动到格子(2,3) ，捕捞 4 条鱼。

示例 2：

输入：grid = [[1,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,1]]
输出：1
解释：渔夫可以从格子 (0,0) 或者 (3,3) ，捕捞 1 条鱼。

提示：

• m == grid.length
• n == grid[i].length
• 1 <= m, n <= 10
• 0 <= grid[i][j] <= 10

问题抽象

求 “加权连通分量 / 岛屿问题”，用二维 BFS 或 DFS 或并查集都可以求出所有连通块的最大值，史上最水 Hard 题。

题解一（二维 DFS）

class Solution {

    private val directions = arrayOf(intArrayOf(0, 1), intArrayOf(0, -1), intArrayOf(1, 0), intArrayOf(-1, 0))

    fun findMaxFish(grid: Array<IntArray>): Int {
        var ret = 0
        for (i in 0 until grid.size) {
            for (j in 0 until grid[0].size) {
                ret = Math.max(ret, dfs(grid, i, j))
            }
        }
        return ret
    }

    private fun dfs(grid: Array<IntArray>, i: Int, j: Int): Int {
        if (grid[i][j] <= 0) return 0
        var cur = grid[i][j]
        grid[i][j] = -1
        for (direction in directions) {
            val newI = i + direction[0]
            val newJ = j + direction[1]
            if (newI < 0 || newI >= grid.size || newJ < 0 || newJ >= grid[0].size || grid[newI][newJ] <= 0) continue
            cur += dfs(grid, newI, newJ)
        }
        return cur
    }
}

复杂度分析：

• 时间复杂度：$O(n · m)$ 其中 n 和 m 是 grid 数组的行和列；
• 空间复杂度：$O(n + m)$ 递归栈的最大深度。

题解二（并查集）

附赠一份并查集的解法：

class Solution {

    private val directions = arrayOf(intArrayOf(0, 1), intArrayOf(0, -1), intArrayOf(1, 0), intArrayOf(-1, 0))

    fun findMaxFish(grid: Array<IntArray>): Int {
        val n = grid.size
        val m = grid[0].size
        var ret = 0
        // 并查集
        val helper = UnionFind(grid)
        // 合并
        for (i in 0 until n) {
            for (j in 0 until m) {
                ret = Math.max(ret, grid[i][j])
                if (grid[i][j] <= 0) continue
                for (direction in directions) {
                    val newI = i + direction[0]
                    val newJ = j + direction[1]
                    if (newI < 0 || newI >= grid.size || newJ < 0 || newJ >= grid[0].size || grid[newI][newJ] <= 0) continue
                    ret = Math.max(ret, helper.union(i * m + j, newI * m + newJ))
                }
            }
        }
        // helper.print()
        return ret
    }

    private class UnionFind(private val grid: Array<IntArray>) {

        private val n = grid.size
        private val m = grid[0].size

        // 父节点
        private val parent = IntArray(n * m) { it }
        // 高度
        private val rank = IntArray(n * m)
        // 数值
        private val value = IntArray(n * m)

        init {
            for (i in 0 until n) {
                for (j in 0 until m) {
                    value[i * m + j] = grid[i][j]
                }
            }
        }

        // return 子集的和
        fun union(x: Int, y: Int): Int {
            // 按秩合并
            val parentX = find(x)
            val parentY = find(y)
            if (parentX == parentY) return value[parentY]
            if (rank[parentX] < rank[parentY]) {
                parent[parentX] = parentY
                value[parentY] += value[parentX]
                return value[parentY]
            } else if (rank[parentY] < rank[parentX]) {
                parent[parentY] = parentX
                value[parentX] += value[parentY]
                return value[parentX]
            } else {
                parent[parentY] = parentX
                value[parentX] += value[parentY]
                rank[parentY]++
                return value[parentX]
            }
        }

        fun print() {
            println("parent=${parent.joinToString()}")
            println("rank=${rank.joinToString()}")
            println("value=${value.joinToString()}")
        }

        private fun find(i: Int): Int {
            // 路径压缩
            var x = i
            while (parent[x] != x) {
                parent[x] = parent[parent[x]]
                x = parent[x]
            }
            return x
        }
    }
}

复杂度分析：

• 时间复杂度：$O(n · m)$ 其中 n 和 m 是 grid 数组的行和列；
• 空间复杂度：$O(n + m)$ 递归栈的最大深度。

相似题目：

• 130. 被围绕的区域
• 200. 岛屿数量
• 990. 等式方程的可满足性

推荐阅读：

• 如何使用并查集解决朋友圈问题？

---

Q4. 将数组清空（Hard）

https://leetcode.cn/problems/make-array-empty/

题目描述

给你一个包含若干 互不相同 整数的数组 nums ，你需要执行以下操作 直到数组为空 ：

• 如果数组中第一个元素是当前数组中的 最小值 ，则删除它。
• 否则，将第一个元素移动到数组的 末尾 。

请你返回需要多少个操作使 nums 为空。

示例 1：

输入：nums = [3,4,-1]
输出：5

Operation	Array
1	[4, -1, 3]
2	[-1, 3, 4]
3	[3, 4]
4	[4]
5	[]

示例 2：

输入：nums = [1,2,4,3]
输出：5

Operation	Array
1	[2, 4, 3]
2	[4, 3]
3	[3, 4]
4	[4]
5	[]

示例 3：

输入：nums = [1,2,3]
输出：3

Operation	Array
1	[2, 3]
2	[3]
3	[]

提示：

• 1 <= nums.length <= 105
• 109 <= nums[i] <= 109
• nums 中的元素 互不相同 。

预备知识 - 循环数组

循环数组：将数组尾部元素的后继视为数组首部元素，数组首部元素的前驱视为数组尾部元素。

预备知识 - 树状数组

OI · 树状数组

树状数组也叫二叉索引树（Binary Indexed Tree），是一种支持 “单点修改” 和 “区间查询” 的代码量少的数据结构。相比于线段树来说，树状数组的代码量远远更少，是一种精妙的数据结构。

树状数组核心思想是将数组 [0,x] 的前缀和拆分为不多于 logx 段非重叠的区间，在计算前缀和时只需要合并 logx 段区间信息，而不需要合并 n 个区间信息。同时，在更新单点值时，也仅需要修改 logx 段区间，而不需要（像前缀和数组）那样修改 n 个信息。可以说，树状数组平衡了单点修改和区间和查询的时间复杂度：

• 单点更新 add(index,val)：将序列第 index 位元素增加 val，时间复杂度为 O(lgn)，同时对应于在逻辑树形结构上从小分块节点移动到大分块节点的过程（修改元素会影响大分块节点（子节点）的值）；
• 区间查询 prefixSum(index)：查询前 index 个元素的前缀和，时间复杂度为 O(lgn)，同时对应于在逻辑树形结构上累加区间段的过程。

树状数组

问题结构化

1、概括问题目标

求消除数组的操作次数。

2、分析题目要件

• 观察：在每次操作中，需要观察数组首部元素是否为剩余元素中的最小值。例如序列 [3,2,1] 的首部元素不是最小值；
• 消除：在每次操作中，如果数组首部元素是最小值，则可以消除数组头部元素。例序列 [1,2,3] 在一次操作后变为 [2,3]；
• 移动：在每次操作中，如果数组首部元素不是最小值，则需要将其移动到数组末尾。例如序列 [3,2,1] 在一次操作后变为 [2,1,3]。

3、观察数据特征

• 数据量：测试用例的数据量上界为 10^5，这要求我们实现低于 O(n^2) 时间复杂度的算法才能通过；
• 数据大小：测试用例的数据上下界为 [-10^9, 10^9]，这要求我们考虑大数问题。

4、观察测试用例

以序列 [3,4,-1] 为例，一共操作 5 次：

• [3,4,-1]：-1 是最小值，将 3 和 4 移动到末尾后才能消除 -1，一共操作 3 次；
• [3,4]：3 是最小值，消除 3 操作 1 次；
• [4]：4 是最小值，消除 4 操作 1 次；

5、提高抽象程度

• 序列：线性表是由多个元素组成的序列，除了数组的头部和尾部元素之外，每个元素都有一个前驱元素和后继元素。在将数组首部元素移动到数组末尾时，将改变数组中的部分元素的关系，即原首部元素的前驱变为原尾部元素，原尾部元素的后继变为原首部元素。
• 是否为决策问题：由于每次操作的行为是固定的，因此这道题只是纯粹的模拟问题，并不是决策问题。

6、具体化解决手段

消除操作需要按照元素值从小到大的顺序删除，那么如何判断数组首部元素是否为最小值？

• 手段 1（暴力枚举）：枚举数组剩余元素，判断首部元素是否为最小值，单次判断的时间复杂度是 O(n)；
• 手段 2（排序）：对原始数组做预处理排序，由于原始数组的元素顺序信息在本问题中是至关重要的，所以不能对原始数组做原地排序，需要借助辅助数据结构，例如索引数组、最小堆，单次判断的均摊时间复杂度是 O(1)。

如何表示元素的移动操作：

• 手段 1（数组）：使用数组块状复制 Arrays.copy()，单次操作的时间复杂度是 O(n)；
• 手段 2（双向链表）：将原始数组转换为双向链表，操作链表首尾元素的时间复杂度是 O(1)，但会消耗更多空间；

如何解决问题：

• 手段 1（模拟）：模拟消除和移动操作，直到数组为空。在最坏情况下（降序数组）需要操作 n^2 次，因此无论如何都是无法满足题目的数据量要求；

至此，问题陷入瓶颈。

解决方法是重复「分析问题要件」-「具体化解决手段」的过程，枚举掌握的算法、数据结构和 Tricks 寻找突破口：

表示元素的移动操作的新手段：

• 手段 3（循环数组）：将原数组视为循环数组，数组尾部元素的后继是数组首部元素，数组首部元素的前驱是数组尾部元素，不再需要实际性的移动操作。

解决问题的新手段：

• 手段 2（计数）：观察测试用例发现，消除每个元素的操作次数取决于该元素的前驱中未被消除的元素个数，例如序列 [3,4,-1] 中 -1 前有 2 个元素未被删除，所以需要 2 次操作移动 3 和 4，再增加一次操作消除 -1。那么，我们可以定义 rangeSum(i,j) 表示区间 [i,j] 中未被删除的元素个数，每次消除操作只需要查询上一次的消除位置（上一个最小值）与当前的消除位置（当前的最小值）中间有多少个数字未被消除 rangeSum(上一个最小值位置, 当前的最小值位置)，这个区间和就是消除当前元素需要的操作次数。

区分上次位置与当前位置的前后关系，需要分类讨论：

• id < preId：消除次数 = rangeSum(id, preId)
• id > preId：消除次数 = rangeSum(-1, id) + rangeSum(preId,n - 1)

如何实现手段 2（计数）：

在代码实现上，涉及到「区间求和」和「单点更新」可以用线段数和树状数组实现。树状数组的代码量远比线段树少，所以我们选择后者。

示意图

答疑：

• 消除每个元素的操作次数不用考虑前驱元素中小于当前元素的元素吗？

由于消除是按照元素值从小到大的顺序消除的，所以未被消除的元素一定比当前元素大，所以我们不强调元素大小关系。

题解一（树状数组 + 索引数组）

• 使用「树状数组」的手段解决区间和查询和单点更新问题，注意树状数组是 base 1 的；
• 使用「索引数组」的手段解决排序 / 最小值问题。

class Solution {
    fun countOperationsToEmptyArray(nums: IntArray): Long {
        val n = nums.size
        var ret = 0L
        // 索引数组
        val ids = Array<Int>(n) { it }
        // 排序
        Arrays.sort(ids) { i1, i2 ->
            // 考虑大数问题
            // nums[i1] - nums[i2] x
            if (nums[i1] < nums[i2]) -1 else 1
        }
        // 树状数组
        val bst = BST(n)
        // 上一个被删除的索引
        var preId = -1
        // 遍历索引
        for (id in ids) {
            // 区间和
            if (id > preId) {
                ret += bst.rangeSum(preId, id)
                // println("id=$id, ${bst.rangeSum(preId, id)}")
            } else {
                ret += bst.rangeSum(-1, id) + bst.rangeSum(preId, n - 1)
                // println("id=$id, ${bst.rangeSum(-1,id)} + ${bst.rangeSum(preId, n - 1)}")
            }
            // 单点更新
            bst.dec(id)
            preId = id
        }
        return ret
    }

    // 树状数组
    private class BST(private val n: Int) {

        // base 1
        private val data = IntArray(n + 1)

        init {
            // O(nlgn) 建树
            // for (i in 0 .. n) {
            //     update(i, 1)
            // }
            // O(n) 建树
            for (i in 1 .. n) {
                data[i] += 1
                val parent = i + lowbit(i)
                if (parent <= n) data[parent] += data[i]
            }
        }

        fun rangeSum(i1: Int, i2: Int): Int {
            return preSum(i2 + 1) - preSum(i1 + 1)
        }

        fun dec(i: Int) {
            update(i + 1, -1)
        }

        private fun preSum(i: Int): Int {
            var x = i
            var sum = 0
            while (x > 0) {
                sum += data[x]
                x -= lowbit(x)
            }
            return sum
        }

        private fun update(i: Int, delta: Int) {
            var x = i
            while (x <= n) {
                data[x] += delta
                x += lowbit(x)
            }
        }

        private fun lowbit(x: Int) = x and (-x)
    }
}

复杂度分析：

• 时间复杂度：$O(nlgn)$ 其中 n 是 nums 数组的长度，排序 $O(nlgn)$、树状数组建树 $O(n)$、单次消除操作的区间和查询和单点更新的时间为 $O(lgn)$；
• 空间复杂度：$O(n)$ 索引数组空间 + 树状数组空间。

题解二（树状数组 + 最小堆）

附赠一份最小堆排序的代码：

• 使用「树状数组」的手段解决区间和查询和单点更新问题，注意树状数组是 base 1 的；
• 使用「最小堆」的手段解决排序 / 最小值问题。

class Solution {
    fun countOperationsToEmptyArray(nums: IntArray): Long {
        val n = nums.size
        var ret = 0L
        // 最小堆
        val ids = PriorityQueue<Int>() { i1, i2 ->
            if (nums[i1] < nums[i2]) -1 else 1
        }
        for (id in 0 until n) {
            ids.offer(id)
        }
        // 树状数组
        val bst = BST(n)
        // 上一个被删除的索引
        var preId = -1
        // 遍历索引
        while (!ids.isEmpty()) {
            val id = ids.poll()
            // 区间和
            if (id > preId) {
                ret += bst.rangeSum(preId, id)
            } else {
                ret += bst.rangeSum(-1, id) + bst.rangeSum(preId, n - 1)
            }
            // 单点更新
            bst.dec(id)
            preId = id
        }
        return ret
    }
}

复杂度分析：

• 时间复杂度：$O(nlgn)$ 其中 n 是 nums 数组的长度，堆排序 $O(nlgn)$、树状数组建树 $O(n)$、单次消除操作的区间和查询和单点更新的时间为 $O(lgn)$；
• 空间复杂度：$O(n)$ 堆空间 + 树状数组空间。

相似题目：

• 315. 计算右侧小于当前元素的个数
• 1040. 移动石子直到连续 II

---

往期回顾

• LeetCode 单周赛第 342 场 · 把问题学复杂，再学简单
• LeetCode 单周赛第 341 场· 难度上来了，图论的问题好多啊！
• LeetCode 双周赛第 102 场· 这次又是最短路。
• LeetCode 双周赛第 101 场 · 是时候做出改变了！