LeetCode 1488. 避免洪水泛滥

2024-07-11 HGNULB 约 14 分钟

文章信息 6 个分类 5 个标签

分类

题目描述

✅ 1488. 避免洪水泛滥

题意分析

给定数组 rains，rains[i] > 0 表示第 $i$ 天在编号为 rains[i] 的湖上下雨、该湖被灌满；rains[i] == 0 表示第 $i$ 天是晴天，可以任选一个已满的湖抽干（也可以什么都不做）。若某天下雨时目标湖已经是满的，就发生洪水。要求构造一个长度相同的答案数组：下雨天填 $-1$，晴天填当天抽干的湖编号（不抽就填任意正整数，通常填 $1$）；无法避免洪水则返回空数组。

把约束翻译一遍：同一个湖若在第 $p$ 天和第 $q$ 天（$p < q$）各下一次雨，那么必须在 $(p, q)$ 这个开区间内找到一个晴天专门抽它。一个晴天只能抽一个湖，所以晴天是稀缺资源，多个湖会争抢同一批晴天——这就是全题唯一的冲突点。

数据规模 $n \le 10^5$ 说明允许 $O(n \log n)$。「必须在某个区间内选一个还没被占用的位置」这种描述天然指向一个支持「查询大于某位置的最小可用元素」并支持删除的动态结构；如果只需要静态查询，前缀和就够了，但晴天被用掉后要从候选里移除，所以结构必须可变。

边界有三处：全是晴天时答案可以随便填正整数，绝不能填 $0$（题目要求晴天位置填的数在 $[1, 10^9]$ 内）；某个湖只下一次雨则不需要为它安排晴天；两次下雨之间一个晴天都没有时立即判定无解。

解法：哈希记录湖泊 + 树状数组分配晴天

核心思路

暴力做法是：扫描到第二次下同一个湖的雨时，回头在 $(pre, i)$ 区间里线性找一个还没被占用的晴天。单次回溯最坏 $O(n)$，总代价 $O(n^2)$，$n = 10^5$ 时是 $10^{10}$，不可行。瓶颈在于「找区间内第一个未占用位置」这个操作被反复地线性重做。

在优化数据结构之前，先要确立贪心策略，否则优化的是错误的目标。当一个湖必须在 $(pre, i)$ 内被抽干时，应该选择该区间内最早的那个可用晴天。交换论证：设最优方案给这个湖分配的是晴天 $b$，而最早可用的是 $a < b$，那么把这个湖改分配到 $a$；若 $a$ 原本被另一个湖 $L'$ 占用，则 $L'$ 的可用区间必然也覆盖 $a$，而 $b$ 在 $a$ 之后且当前空出来了，$L'$ 改用 $b$ 是否合法要看 $b$ 是否仍在 $L'$ 的区间内——由于我们是按时间顺序处理的，$L'$ 的右端点在 $i$ 之前，$b < i$ 但 $b$ 未必小于 $L'$ 的右端点。所以更严谨的说法是：按下雨事件的时间顺序处理，每次取当前区间内最早的可用晴天，永远不会比其他选择更差——因为把早的晴天留着只可能被后来右端点更大的区间使用，而后来的区间左端点也更大，早的晴天对它们反而可能已经失效。直觉上就是「晚的晴天更通用，要留给后面」。

有了贪心策略，剩下的就是把「区间内第一个可用晴天」做快。这里用树状数组维护一个 01 序列：dryDays 的第 $t$ 位为 $1$ 表示第 $t$ 天是晴天且尚未被分配，为 $0$ 表示不是晴天或已被用掉。它天然支持两个操作——前缀和 sum(t) 给出前 $t$ 天里可用晴天的个数，单点 add(t, ±1) 完成「新增一个可用晴天」与「消耗一个晴天」。

关键一步是把「找区间内第一个可用晴天」转成「找第 $r$ 个可用晴天」这个排名查询：设 usedBefore = sum(pre + 1) 是第 $pre$ 天及之前还剩多少可用晴天，那么 $(pre, i)$ 内第一个可用晴天就是全局第 usedBefore + 1 个可用晴天。树状数组上做排名查询可以在 $O(\log n)$ 内完成，方法是从最高位的 $2$ 的幂开始逐位试探下降（findByOrder），这比「二分套前缀和」的 $O(\log^2 n)$ 更快。

无解判定也顺势得到：若 total() == usedBefore，说明第 $pre$ 天之后一个可用晴天都没有，直接返回空数组。注意这里不需要额外检查「找到的晴天是否小于 $i$」——因为第 $i$ 天本身是雨天（不是可用晴天），而 $i$ 之后的晴天此刻还没被加入树状数组（我们是边扫描边插入的），所以排名查到的结果必然落在 $(pre, i)$ 内。这个「只把已扫描过的晴天入表」的时序安排，免掉了右端点的判断。

解题步骤

准备三样东西：结果数组 answer、哈希表 lastRain（湖编号到它上次下雨的天数）、树状数组 dryDays（长度 $n$ 的 01 计数）。哈希表存「上次」而不是全部历史，是因为只有相邻两次下雨之间才需要安排抽水，更早的雨已经被此前的抽水解决了。

从左到右扫描每一天，遇到晴天（lake == 0）时先把 answer[i] 填 $1$，再执行 dryDays.add(i + 1, 1)。填 $1$ 是占位默认值，代表「这天暂时不抽任何湖」，后续若被某个湖认领会被覆盖；填 $1$ 而不是 $0$ 是因为题目要求这个数必须是正的合法湖编号。树状数组下标从 $1$ 起，所以第 $i$ 天对应位置 $i + 1$。

遇到雨天时先把 answer[i] 填 $-1$，这是题目对雨天位置的硬性规定，与是否需要抽水无关。

若 lastRain 中已有这个湖，取出上次下雨的天数 pre，计算 usedBefore = dryDays.sum(pre + 1)。这个前缀和的含义是「第 $pre$ 天（含）之前还剩几个没被用掉的晴天」；由于第 $pre$ 天本身是雨天，位置 $pre+1$ 上的值是 $0$，所以包不包含它都不影响结果。

若 dryDays.total() == usedBefore 立即返回空数组。两者相等意味着所有可用晴天都在 $pre$ 之前，$(pre, i)$ 区间里一个都没有，这个湖必然溢出，整个方案无解，不必再扫描后面。

否则用 findByOrder(usedBefore + 1) 取出第 usedBefore + 1 个可用晴天的位置 dryPos，这正是 $pre$ 之后最早的那个。排名查询用树状数组自身的树形结构逐位下降，每一步用 tree[next] < order 判断是否可以往右跨越一整块，跨越后从 order 里扣掉这一块的计数。

把 answer[dryPos - 1] 改写成 lake，并执行 dryDays.add(dryPos, -1)。前者把那天的默认占位值替换成真正要抽的湖；后者把这个晴天从可用集合中删除，保证它不会被第二个湖重复认领——这正是「一个晴天只能抽一个湖」的落实。

无论是否需要抽水，最后都执行 lastRain.put(lake, i)，把这个湖的「上次下雨」推进到今天，供下一次相遇时使用。

以 rains = [1, 2, 0, 0, 2, 1] 走一遍（树状数组位置用 1-indexed 天数表示）：

$i = 0$：lake = 1，answer[0] = -1，lastRain 中无湖 1，直接记 lastRain[1] = 0。
$i = 1$：lake = 2，answer[1] = -1，记 lastRain[2] = 1。
$i = 2$：晴天，answer[2] = 1，树状数组位置 $3$ 置 $1$。可用晴天集合 ${3}$。
$i = 3$：晴天，answer[3] = 1，位置 $4$ 置 $1$。可用集合 ${3, 4}$。
$i = 4$：lake = 2，answer[4] = -1。pre = 1，usedBefore = sum(2) = 0（位置 1、2 上没有可用晴天）。total() = 2 \ne 0，有解。findByOrder(1) 返回位置 $3$，即第 2 天。于是 answer[3 - 1] = answer[2] = 2，并把位置 $3$ 清零。可用集合变为 ${4}$。更新 lastRain[2] = 4。
$i = 5$：lake = 1，answer[5] = -1。pre = 0，usedBefore = sum(1) = 0。total() = 1 \ne 0。findByOrder(1) 返回位置 $4$，即第 3 天。answer[4 - 1] = answer[3] = 1，位置 $4$ 清零。更新 lastRain[1] = 5。
最终 answer = [-1, -1, 2, 1, -1, -1]。

这个用例正好体现了贪心的必要性：$i = 4$ 时湖 2 的可选晴天有第 2 天和第 3 天，若贪心地选了更晚的第 3 天，那么 $i = 5$ 时湖 1 需要 $(0, 5)$ 内的晴天，只剩第 2 天可用——恰好也能成。但把区间收紧一点，比如 rains = [1, 0, 2, 0, 2, 1] 这类构造中，选晚的晴天就会让后面某个左端点更大的湖无天可用。规则是统一的：永远取最早可用的那个。

代码实现

class Solution {
    public int[] avoidFlood(int[] rains) {
        int n = rains.length;
        int[] answer = new int[n];
        Map<Integer, Integer> lastRain = new HashMap<>();
        Fenwick dryDays = new Fenwick(n);

        for (int i = 0; i < n; i++) {
            int lake = rains[i];
            if (lake == 0) {
                answer[i] = 1;
                dryDays.add(i + 1, 1);
                continue;
            }

            answer[i] = -1;
            if (lastRain.containsKey(lake)) {
                int pre = lastRain.get(lake);
                int usedBefore = dryDays.sum(pre + 1);
                if (dryDays.total() == usedBefore) {
                    return new int[0];
                }

                int dryPos = dryDays.findByOrder(usedBefore + 1);
                answer[dryPos - 1] = lake;
                dryDays.add(dryPos, -1);
            }
            lastRain.put(lake, i);
        }

        return answer;
    }

    private static class Fenwick {
        private final int[] tree;

        Fenwick(int n) {
            tree = new int[n + 1];
        }

        void add(int idx, int delta) {
            while (idx < tree.length) {
                tree[idx] += delta;
                idx += idx & -idx;
            }
        }

        int sum(int idx) {
            int answer = 0;
            while (idx > 0) {
                answer += tree[idx];
                idx -= idx & -idx;
            }
            return answer;
        }

        int total() {
            return sum(tree.length - 1);
        }

        int findByOrder(int order) {
            int idx = 0;
            int step = 1;
            while (step < tree.length) {
                step <<= 1;
            }

            while (step > 0) {
                int next = idx + step;
                if (next < tree.length && tree[next] < order) {
                    idx = next;
                    order -= tree[next];
                }
                step >>= 1;
            }

            return idx + 1;
        }
    }
}

func avoidFlood(rains []int) []int {
    n := len(rains)
    answer := make([]int, n)
    lastRain := make(map[int]int)
    dryDays := newFenwick(n)

    for i, lake := range rains {
        if lake == 0 {
            answer[i] = 1
            dryDays.add(i+1, 1)
            continue
        }

        answer[i] = -1
        if pre, ok := lastRain[lake]; ok {
            usedBefore := dryDays.sum(pre + 1)
            if dryDays.total() == usedBefore {
                return []int{}
            }

            dryPos := dryDays.findByOrder(usedBefore + 1)
            answer[dryPos-1] = lake
            dryDays.add(dryPos, -1)
        }
        lastRain[lake] = i
    }

    return answer
}

type fenwick struct {
    tree []int
}

func newFenwick(n int) *fenwick {
    return &fenwick{tree: make([]int, n+1)}
}

func (f *fenwick) add(idx int, delta int) {
    for idx < len(f.tree) {
        f.tree[idx] += delta
        idx += idx & -idx
    }
}

func (f *fenwick) sum(idx int) int {
    answer := 0
    for idx > 0 {
        answer += f.tree[idx]
        idx -= idx & -idx
    }
    return answer
}

func (f *fenwick) total() int {
    return f.sum(len(f.tree) - 1)
}

func (f *fenwick) findByOrder(order int) int {
    idx := 0
    step := 1
    for step < len(f.tree) {
        step <<= 1
    }

    for step > 0 {
        next := idx + step
        if next < len(f.tree) && f.tree[next] < order {
            idx = next
            order -= f.tree[next]
        }
        step >>= 1
    }

    return idx + 1
}

复杂度分析

时间复杂度：$O(n \log n)$。外层扫描 $n$ 天；每个晴天做一次树状数组单点更新 $O(\log n)$；每个「第二次及以后下雨」的雨天做一次前缀和 $O(\log n)$、一次排名查询 $O(\log n)$、一次单点更新 $O(\log n)$。哈希表的插入与查询均摊 $O(1)$。总计 $O(n \log n)$，$n = 10^5$ 时约 $1.7 \times 10^6$ 次基本操作。

空间复杂度：$O(n)$。结果数组 $n$ 个整数，树状数组 $n + 1$ 个整数，哈希表最多存 $n$ 个不同的湖编号。没有递归，栈空间是常数。

关键点总结

遇到「多个需求争抢同一批稀缺资源，每个需求有一个可用区间」的题，先确立贪心顺序：按需求的右端点（截止时间）从早到晚处理，每个需求取区间内最早的可用资源。留下的晚资源对后续左端点更大的需求更通用，这条直觉几乎在所有区间分配题里都成立。

贪心策略确立之后，才是选数据结构。「找大于某位置的最小可用元素 + 删除」这组操作，有序集合（TreeSet.ceiling）最直白，树状数组维护 01 序列 + 排名查询次之，二者都是 $O(\log n)$；树状数组的优势是常数小、无对象开销，且排名查询可以做到单次 $O(\log n)$ 而非 $O(\log^2 n)$。

树状数组上「找第 $r$ 个 1」的逐位下降写法要记牢：从最高的 $2$ 的幂开始，若 tree[idx + step] < order 就整块跨过去并从 order 里扣掉。它把二分与前缀和融进了同一次遍历。

边扫描边插入可用资源，能天然保证查到的资源不会越过当前时刻，省掉右端点校验。设计算法时刻意安排「什么时候把元素放进结构」，往往比事后加判断更干净。

面试视角：面试官会依次考三件事——能不能说清「两次下雨之间必须有一个专属晴天」这个建模；能不能论证「取最早可用晴天」的贪心正确性（要会用交换论证，并解释「晚的晴天更通用」）；能不能给出合适的数据结构并说明复杂度。Java 现场手写建议直接用 TreeSet 加 ceiling，把时间花在贪心论证上；被追问「不用有序集合怎么办」时，再抛出树状数组维护 01 序列做排名查询这条路。

易错点总结

晴天默认值填 $0$：rains = [0] 时返回 [0]，而题目要求晴天位置填 $[1, 10^9]$ 内的正整数，判定不通过；应填 $1$。

雨天位置忘记填 $-1$：rains = [1, 2] 时返回 [0, 0]，正确答案是 [-1, -1]。

贪心时取区间内最晚的可用晴天：构造两个湖的截止时间交错的用例（如湖 A 的区间是 $(0, 3)$、湖 B 的区间是 $(2, 5)$，可用晴天在第 1、4 天）时，湖 A 抢走第 4 天会让湖 B 无天可用，返回空数组，而实际有解。

抽干晴天后忘记从可用集合中删除：rains = [1, 2, 0, 1, 2] 时第 2 天会被同时分配给湖 1 和湖 2，answer[2] 被覆盖成后写入的那个，另一个湖实际溢出，输出是一个非法方案。

lastRain 记录的是首次下雨而非最近一次：rains = [1, 0, 1, 0, 1] 中第三次下雨时会用 pre = 0 去查区间，可能选中第 1 天那个已经被第二次下雨用掉的晴天，导致重复占用判断错乱。

无解时继续扫描而不立即返回：rains = [1, 1, 0] 时第 $i = 1$ 天已经确定溢出，若继续走完循环并返回 answer，会输出一个看似合法实则洪水的方案，正确应返回空数组。

前缀和的下标没有从 0-indexed 天数转成 1-indexed 位置：把 dryDays.add(i, 1) 写成不加一，第 $0$ 天的晴天会更新到位置 $0$，而树状数组的 add 在 idx = 0 时 idx & -idx 为 $0$，循环永不前进，直接死循环。

把 answer[dryPos] 而非 answer[dryPos - 1] 赋值：rains = [1, 0, 1] 中 dryPos = 2 对应第 1 天，写成 answer[2] 会把雨天位置的 $-1$ 覆盖掉，输出 [-1, 1, 1]，正确是 [-1, 1, -1]。

判无解写成 usedBefore == 0：rains = [1, 0, 2, 1] 中处理湖 1 时 usedBefore = 0，但第 1 天的晴天确实可用，会被误判成无解返回空数组，正确答案存在。

findByOrder 里 step 的初值没有取到不小于表长的 $2$ 的幂：$n = 5$ 时若从 step = 4 开始而表长为 $6$，最高位块被跳过，排名查询返回偏小的位置，分配到一个已被用掉或不存在的晴天。

相似题目

题目	难度	考察点
630. 课程表 III	困难	同为截止时间约束下的资源分配，用大顶堆反悔贪心而非有序集合查找
253. 会议室 II	中等	也是稀缺资源复用，但只需数并发峰值，小顶堆维护最早结束时间即可
621. 任务调度器	中等	冷却期约束下的时间片分配，靠最高频任务的桶结构直接算出答案
502. IPO	困难	双堆贪心，按资金门槛解锁候选再取最大收益，考的是「何时把元素放进结构」
315. 计算右侧小于当前元素的个数	困难	同样用树状数组维护 01 计数，但做的是逆序对统计而非可用位置分配
307. 区域和检索 - 数组可修改	中等	树状数组的标准模板题，只需单点改与前缀和，不涉及排名查询
452. 用最少数量的箭引爆气球	中等	区间贪心的另一形态，按右端点排序后一次覆盖尽可能多的区间
220. 存在重复元素 III	困难	同为「在有序结构里找不小于某值的最小元素」，但结构随滑动窗口动态增删