1488. 避免洪水泛滥

题目

你的国家有无数个湖泊,所有湖泊一开始都是空的。当第 n 个湖泊下雨前是空的,那么它就会装满水。如果第 n 个湖泊下雨前是 满的 ,这个湖泊会发生 洪水 。你的目标是避免任意一个湖泊发生洪水。

给你一个整数数组 rains ,其中:

  • rains[i] > 0 表示第 i 天时,第 rains[i] 个湖泊会下雨。
  • rains[i] == 0 表示第 i 天没有湖泊会下雨,你可以选择 一个 湖泊并 抽干 这个湖泊的水。

请返回一个数组 ans ,满足:

  • ans.length == rains.length
  • 如果 rains[i] > 0 ,那么ans[i] == -1 。
  • 如果 rains[i] == 0 ,ans[i] 是你第 i 天选择抽干的湖泊。

如果有多种可行解,请返回它们中的 任意一个 。如果没办法阻止洪水,请返回一个 空的数组 。

请注意,如果你选择抽干一个装满水的湖泊,它会变成一个空的湖泊。但如果你选择抽干一个空的湖泊,那么将无事发生。

示例 1:

输入:rains = [1,2,3,4] 输出:[-1,-1,-1,-1] 解释:第一天后,装满水的湖泊包括 [1] 第二天后,装满水的湖泊包括 [1,2] 第三天后,装满水的湖泊包括 [1,2,3] 第四天后,装满水的湖泊包括 [1,2,3,4] 没有哪一天你可以抽干任何湖泊的水,也没有湖泊会发生洪水。

示例 2:

输入:rains = [1,2,0,0,2,1] 输出:[-1,-1,2,1,-1,-1] 解释:第一天后,装满水的湖泊包括 [1] 第二天后,装满水的湖泊包括 [1,2] 第三天后,我们抽干湖泊 2 。所以剩下装满水的湖泊包括 [1] 第四天后,我们抽干湖泊 1 。所以暂时没有装满水的湖泊了。 第五天后,装满水的湖泊包括 [2]。 第六天后,装满水的湖泊包括 [1,2]。 可以看出,这个方案下不会有洪水发生。同时, [-1,-1,1,2,-1,-1] 也是另一个可行的没有洪水的方案。

示例 3:

输入:rains = [1,2,0,1,2] 输出:[] 解释:第二天后,装满水的湖泊包括 [1,2]。我们可以在第三天抽干一个湖泊的水。 但第三天后,湖泊 1 和 2 都会再次下雨,所以不管我们第三天抽干哪个湖泊的水,另一个湖泊都会发生洪水。

提示:

  • 1 <= rains.length <= 10^5
  • 0 <= rains[i] <= 10^9

解题思路

1.被动响应,延迟决策

  • 核心思想: “懒惰决策”。在不下雨的日子,我们什么都不做,只是把这个“可以抽水的机会”记录下来。只有在真正要发生洪水的那一刻,才被迫回头去利用之前存储的机会。
  • 决策时机: 当天下雨,并且下的这个湖已经满了,即将发生洪水。
  • 决策方法:
    1. 查看所有我们积攒下来的“可以抽水的日子”。
    2. 在这些日子中,找到一个在上一次下雨之后、本次下雨之前的。
    3. 为了给未来留下更多(更靠后)的机会,我们贪心地选择其中最早的那个抽水日来解决这次危机。
  • 所需数据结构:
    1. 一个字典/哈希表 (map): 记录已满的湖以及它上一次下雨的日期。
    2. 一个有序的列表 (sliceBalanced BST): 存储所有可以抽水的日子的索引。
  • 复杂度与瓶颈:
    • 如果用普通切片存储抽水日,查找 ($O(logN)$) 很快,但删除 ($O(N)$) 很慢,导致整体最坏复杂度为 $O(N^2)$。
    • 如果用平衡二叉搜索树,查找和删除都可以做到 $O(logN)$,使整体复杂度达到最优的 $O(NlogN)$。

2.主动规划的贪心

  • 核心思想: “防患未然”。我们首先通过预处理“看穿”未来,知道所有湖泊的下雨计划。然后在每一个不下雨的日子,我们主动出击,解决掉那个最迫在眉睫的潜在危机。
  • 决策时机: 当天不下雨,我们拥有一次抽水的机会。
  • 决策方法:
    1. 查看所有当前已满的湖。
    2. 对于每一个已满的湖,查询它的“天气预报”,找到它下一次会在哪天再次下雨。
    3. 这些“下一次下雨日”代表了未来的一个个潜在危机。我们贪心地选择其中**日期最早(最紧急)**的那个危机,在今天这个不下雨的日子把它解决掉(抽干对应的湖)。
  • 所需数据结构:
    1. 一个字典/哈希表 (map): 作为“天气预报”,记录每个湖未来所有的下雨日期。
    2. 一个集合 (set): 记录当前哪些湖是满的。
    3. 一个最小堆 (Min-Heap): 作为“紧急任务板”,存储所有已满湖泊的下一次下雨日期,并自动将日期最早的顶在最上面。
  • 复杂度与瓶颈:
    • 这种方法天生就需要高效地找到“最小值”,所以最小堆是完美的选择。
    • 堆的插入和删除操作都是 $O(logN)$ 的。
    • 因此,这种思路的实现很自然地就是 $O(NlogN)$ 的最优解,没有之前那种 $O(N^2)$ 的陷阱。

具体代码

被动思路

func avoidFlood(rains []int) []int {
    n := len(rains)
    ans := make([]int, n)
    // 默认给所有不下雨的日子填充一个抽干湖1的选项
    // 如果后续需要用来避免洪水,这个值会被覆盖
    for i := 0; i < n; i++ {
        ans[i] = 1
    }

    // key: 湖的编号, value: 上次下雨的日期
    fullLakes := make(map[int]int)
    // 用来存储所有可以抽水的日子的索引 (天然有序)
    dryDays := make([]int, 0)

    for day, lake := range rains {
        if lake == 0 {
            // 今天不下雨,是一个可以抽水的机会,存起来
            dryDays = append(dryDays, day)
        } else {
            // 今天下雨,不能抽水
            ans[day] = -1

            if prevRainDay, ok := fullLakes[lake]; ok {
                // 这个湖之前已经满了,即将洪水
                // 我们需要在 dryDays 中找到一个比 prevRainDay 大的、最小的日子来抽水
                
                // 使用二分查找在 dryDays 中寻找第一个 > prevRainDay 的索引
                // sort.Search 会返回第一个满足条件的元素的索引
                idx := sort.Search(len(dryDays), func(i int) bool {
                    return dryDays[i] > prevRainDay
                })

                if idx == len(dryDays) {
                    // 没有找到任何一个在上次下雨之后的抽水日
                    // 无法避免洪水
                    return []int{}
                }

                // 找到了最合适的抽水日
                drainDay := dryDays[idx]
                ans[drainDay] = lake

                // 这个抽水日已经被用掉了,从 dryDays 中移除
                // Go中切片删除元素的常用写法
                dryDays = append(dryDays[:idx], dryDays[idx+1:]...)

            }
            
            // 更新这个湖最后一次下雨的日期
            fullLakes[lake] = day
        }
    }
    
    return ans
}

主动思路

package main

import (
	"container/heap"
	"fmt"
)

// --- 第一部分: 最小堆的实现 (紧急任务板) ---
// 这是一个标准模板,可以用于任何需要整数最小堆的场景。

// IntHeap 是一个整数最小堆
type IntHeap []int

// Len, Less, Swap 是为了满足 sort.Interface 接口
func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // h[i] < h[j] 表示最小堆
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

// Push 和 Pop 是为了满足 heap.Interface 接口
func (h *IntHeap) Push(x any) {
	// Push 和 Pop 使用指针接收者,因为它们会修改切片的长度。
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// --- 第二部分: avoidFlood 主函数 ---

func avoidFlood(rains []int) []int {
	n := len(rains)
	ans := make([]int, n)

	// 1. 预处理,生成 "天气预报表"
	// key: 湖编号, value: 该湖所有下雨日期的切片
	lakeToRainDays := make(map[int][]int)
	for day, lake := range rains {
		if lake > 0 {
			lakeToRainDays[lake] = append(lakeToRainDays[lake], day)
		}
	}

	// 2. 初始化 "已满的湖" 集合和 "紧急任务板" 最小堆
	fullLakes := make(map[int]struct{}) // Go中实现集合的标准方式
	pq := &IntHeap{}                     // 优先队列 (最小堆)

	// 3. 开始按天处理
	for day, lake := range rains {
		if lake > 0 { // 今天是下雨天
			if _, ok := fullLakes[lake]; ok {
				// 这个湖已经满了,无法避免洪水
				return []int{}
			}

			ans[day] = -1
			fullLakes[lake] = struct{}{} // 标记为已满

			// 从"天气预报表"中消耗掉今天的降雨记录
			// 注意:这里我们不需要真的删除,因为循环是按顺序的,
			// 我们只需要在预报表中找到下一次的降雨即可。
			// 但为了和Python版本逻辑保持完全一致,我们同样“消耗”它。
			lakeToRainDays[lake] = lakeToRainDays[lake][1:]

			// 如果这个湖未来还会下雨,将下一次的日期加入紧急任务板
			if len(lakeToRainDays[lake]) > 0 {
				nextRainDay := lakeToRainDays[lake][0]
				heap.Push(pq, nextRainDay)
			}
		} else { // 今天是晴天,可以抽水
			if pq.Len() == 0 {
				// 任务板是空的,没有紧急任务
				// 随便抽一个湖,这里用 1 作为占位符
				ans[day] = 1
			} else {
				// 有紧急任务!处理最紧急的那个
				dayOfNextRain := heap.Pop(pq).(int)
				lakeToDry := rains[dayOfNextRain]

				ans[day] = lakeToDry // 决策:今天抽干这个最危险的湖

				delete(fullLakes, lakeToDry) // 将湖从"已满"集合中移除
			}
		}
	}

	return ans
}