1488. Avoid Flood in The City

MediumGreedyArrayHash TableBinary SearchHeap (Priority Queue)

Problem Description

You have an infinite number of lakes numbered 1, 2, 3, and so on. Initially, all lakes are empty. You're given an array rains that represents weather conditions over consecutive days:

When rains[i] > 0, it means it will rain over lake number rains[i] on day i, filling it with water
When rains[i] == 0, it means day i is sunny, and you can choose to dry exactly one lake (making it empty)

The key constraint is that if it rains over a lake that's already full of water, a flood occurs. Your task is to prevent all floods by strategically choosing which lakes to dry on sunny days.

You need to return an array ans of the same length as rains where:

ans[i] = -1 when rains[i] > 0 (rainy day, no action needed)
ans[i] = the lake number you choose to dry when rains[i] == 0 (sunny day)

If multiple valid solutions exist, you can return any of them. If it's impossible to prevent floods, return an empty array.

Note that drying an already empty lake is allowed but has no effect - it remains empty.

For example, if rains = [1, 2, 0, 0, 2, 1]:

Day 0: Lake 1 fills with water
Day 1: Lake 2 fills with water
Day 2: Sunny day - you could dry lake 2
Day 3: Sunny day - you could dry lake 1
Day 4: Lake 2 fills again (was dried on day 2)
Day 5: Lake 1 fills again (was dried on day 3)

A valid answer would be [-1, -1, 2, 1, -1, -1].

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

The key insight is that we need to be strategic about when to dry lakes. We can't just dry any lake on a sunny day - we need to dry lakes that will flood in the future.

Think about it this way: if a lake gets rained on twice, we must dry it between those two rain events. The challenge is that we might not have a sunny day exactly when we need it. So when we encounter a lake that's about to be rained on for the second time, we need to look back and find a sunny day between the first and second rain to dry it.

But here's the crucial observation: if we have multiple sunny days available between two rains on the same lake, which one should we use? We should use the earliest available sunny day after the first rain. Why? Because this preserves our flexibility - we keep later sunny days available for other lakes that might need them.

This leads us to a greedy strategy:

Keep track of all available sunny days
Remember when each lake was last rained on
When a lake is about to be rained on again, find the first available sunny day after its previous rain and use it to dry that lake

The reason we need binary search is efficiency - when looking for "the first sunny day after a specific date", we can use binary search on our collection of available sunny days to find it quickly.

If at any point we can't find a sunny day between two rains on the same lake, it means flooding is inevitable and we return an empty array.

This greedy approach works because using the earliest possible sunny day for each lake maximizes our options for handling future conflicts. Delaying the drying unnecessarily might cause us to miss opportunities to prevent floods on other lakes.

Learn more about Greedy, Binary Search and Heap (Priority Queue) patterns.

Solution Approach

Let's implement the greedy strategy using appropriate data structures:

Data Structures:

sunny: A sorted list (or set) to store all sunny days (when rains[i] == 0). We use a sorted structure for efficient binary search.
rainy: A hash table that maps each lake number to the day it was last rained on.
ans: The answer array initialized with -1 for all positions.

Algorithm Steps:

Initialize the data structures and set all elements in ans to -1.
Traverse through each day in the rains array:

For a rainy day (when rains[i] > 0):
- Check if this lake has been rained on before by looking up rainy[rains[i]].
- If yes, we need to find a sunny day to dry it:
  - Use binary search (bisect_right) to find the first sunny day after the previous rain (rainy[rains[i]]).
  - If no such sunny day exists (index equals length of sunny), flooding is inevitable - return empty array.
  - Otherwise, mark that sunny day to dry this lake: ans[sunny[idx]] = rains[i].
  - Remove that sunny day from our available days: sunny.discard(sunny[idx]).
- Update rainy[rains[i]] = i to record this as the latest rain on this lake.
For a sunny day (when rains[i] == 0):
- Add this day to our collection of available sunny days: sunny.add(i).
- Set a default value ans[i] = 1 (we can dry any lake, defaulting to lake 1).
Return the answer array after processing all days.

Why Binary Search? When we need to find a sunny day to prevent flooding, we specifically need the first sunny day after the previous rain. Using bisect_right on the sorted sunny list gives us exactly this - the smallest index where we can insert the previous rain day while maintaining sorted order, which corresponds to the first sunny day after it.

Time Complexity: O(n log n) where n is the length of the rains array. Each operation on the sorted list takes O(log n) time.

Space Complexity: O(n) for storing the sunny days, rainy days hash table, and answer array.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through the solution with rains = [1, 2, 0, 0, 2, 1]:

Initial Setup:

sunny = [] (sorted list of available sunny days)
rainy = {} (tracks last rain day for each lake)
ans = [-1, -1, -1, -1, -1, -1] (all initialized to -1)

Day 0: rains[0] = 1 (Lake 1 gets rained on)

Lake 1 hasn't been rained on before (not in rainy)
Record this rain: rainy[1] = 0
No action needed, ans[0] remains -1
State: rainy = {1: 0}, sunny = []

Day 1: rains[1] = 2 (Lake 2 gets rained on)

Lake 2 hasn't been rained on before
Record this rain: rainy[2] = 1
No action needed, ans[1] remains -1
State: rainy = {1: 0, 2: 1}, sunny = []

Day 2: rains[2] = 0 (Sunny day)

Add to available sunny days: sunny = [2]
Set default: ans[2] = 1
State: rainy = {1: 0, 2: 1}, sunny = [2]

Day 3: rains[3] = 0 (Sunny day)

Add to available sunny days: sunny = [2, 3]
Set default: ans[3] = 1
State: rainy = {1: 0, 2: 1}, sunny = [2, 3]

Day 4: rains[4] = 2 (Lake 2 gets rained on again)

Lake 2 was previously rained on day 1 (check rainy[2] = 1)
Need a sunny day after day 1 to dry it
Binary search in sunny = [2, 3] for first day > 1
Found day 2 at index 0
Use day 2 to dry lake 2: ans[2] = 2
Remove day 2 from sunny days: sunny = [3]
Update last rain: rainy[2] = 4
ans[4] remains -1
State: rainy = {1: 0, 2: 4}, sunny = [3]

Day 5: rains[5] = 1 (Lake 1 gets rained on again)

Lake 1 was previously rained on day 0 (check rainy[1] = 0)
Need a sunny day after day 0 to dry it
Binary search in sunny = [3] for first day > 0
Found day 3 at index 0
Use day 3 to dry lake 1: ans[3] = 1
Remove day 3 from sunny days: sunny = []
Update last rain: rainy[1] = 5
ans[5] remains -1
State: rainy = {1: 5, 2: 4}, sunny = []

Final Answer: [-1, -1, 2, 1, -1, -1]

This successfully prevents all floods by:

Drying lake 2 on day 2 (before it rains again on day 4)
Drying lake 1 on day 3 (before it rains again on day 5)

Solution Implementation

1from typing import List
2from sortedcontainers import SortedList
3
4class Solution:
5    def avoidFlood(self, rains: List[int]) -> List[int]:
6        n = len(rains)
7        # Initialize result array with -1 (default for rainy days)
8        result = [-1] * n
9      
10        # Keep track of sunny days (when we can dry a lake)
11        sunny_days = SortedList()
12      
13        # Dictionary to store the last day each lake was filled
14        # Key: lake number, Value: day index when it rained on this lake
15        last_rain_day = {}
16      
17        for day, lake in enumerate(rains):
18            if lake > 0:  # Rainy day (lake > 0 means it rains on this lake)
19                # Check if this lake was already full
20                if lake in last_rain_day:
21                    # Find the first sunny day after the last rain on this lake
22                    # We need to dry this lake before it rains again
23                    sunny_day_idx = sunny_days.bisect_right(last_rain_day[lake])
24                  
25                    # If no sunny day available after last rain, flood is inevitable
26                    if sunny_day_idx == len(sunny_days):
27                        return []
28                  
29                    # Use this sunny day to dry the current lake
30                    dry_day = sunny_days[sunny_day_idx]
31                    result[dry_day] = lake
32                  
33                    # Remove the used sunny day from available days
34                    sunny_days.discard(dry_day)
35              
36                # Update the last rain day for this lake
37                last_rain_day[lake] = day
38            else:  # Sunny day (lake == 0)
39                # Add this day to available sunny days
40                sunny_days.add(day)
41                # Default: dry lake 1 (can be any valid lake number)
42                result[day] = 1
43      
44        return result
45

1class Solution {
2    public int[] avoidFlood(int[] rains) {
3        int n = rains.length;
4        int[] result = new int[n];
5        // Initialize all days with -1 (default for rainy days)
6        Arrays.fill(result, -1);
7      
8        // TreeSet to store indices of sunny days (when rains[i] == 0)
9        // Using TreeSet for efficient search of the next available sunny day
10        TreeSet<Integer> sunnyDays = new TreeSet<>();
11      
12        // Map to store the last occurrence index of each lake that was filled
13        // Key: lake number, Value: day index when it last rained on this lake
14        Map<Integer, Integer> lastRainedLakes = new HashMap<>();
15      
16        for (int day = 0; day < n; day++) {
17            int lake = rains[day];
18          
19            if (lake > 0) {
20                // It's raining on this lake
21                if (lastRainedLakes.containsKey(lake)) {
22                    // This lake was already full, we need to dry it before today
23                    // Find the first sunny day after the last time it rained on this lake
24                    Integer dryDay = sunnyDays.higher(lastRainedLakes.get(lake));
25                  
26                    if (dryDay == null) {
27                        // No sunny day available to dry this lake - flooding occurs
28                        return new int[0];
29                    }
30                  
31                    // Use this sunny day to dry the current lake
32                    result[dryDay] = lake;
33                    sunnyDays.remove(dryDay);
34                }
35                // Update the last rained day for this lake
36                lastRainedLakes.put(lake, day);
37            } else {
38                // It's a sunny day, add to available sunny days
39                sunnyDays.add(day);
40                // Temporarily set to 1 (can be any positive number if not used)
41                result[day] = 1;
42            }
43        }
44      
45        return result;
46    }
47}
48

1class Solution {
2public:
3    vector<int> avoidFlood(vector<int>& rains) {
4        int n = rains.size();
5        vector<int> result(n, -1);  // Initialize result array with -1 (default for rainy days)
6      
7        // Set to store indices of sunny days (when rains[i] == 0)
8        set<int> sunnyDays;
9      
10        // Map to store the last occurrence index of each lake that got filled
11        unordered_map<int, int> lastRainedLakes;
12      
13        for (int i = 0; i < n; ++i) {
14            int lake = rains[i];
15          
16            if (lake > 0) {  // Rainy day - lake gets filled
17                // Check if this lake was already filled before
18                if (lastRainedLakes.count(lake)) {
19                    // Find the first sunny day after the last time this lake was filled
20                    auto sunnyDayIt = sunnyDays.upper_bound(lastRainedLakes[lake]);
21                  
22                    // If no sunny day available to dry this lake, flooding occurs
23                    if (sunnyDayIt == sunnyDays.end()) {
24                        return {};  // Return empty array indicating failure
25                    }
26                  
27                    // Use this sunny day to dry the current lake
28                    result[*sunnyDayIt] = lake;
29                    sunnyDays.erase(sunnyDayIt);  // Remove used sunny day
30                }
31              
32                // Update the last occurrence of this lake being filled
33                lastRainedLakes[lake] = i;
34            } else {  // Sunny day - can dry one lake
35                sunnyDays.insert(i);  // Add this sunny day to available days
36                result[i] = 1;  // Temporarily set to 1 (can be any valid lake number)
37            }
38        }
39      
40        return result;
41    }
42};
43

1/**
2 * Solution for "Avoid Flood in The City" problem
3 * Given an array where rains[i] > 0 means lake rains[i] will be filled,
4 * and rains[i] == 0 means we can dry one lake.
5 * Returns an array where -1 indicates rain, and positive values indicate which lake to dry.
6 * Returns empty array if flooding is unavoidable.
7 */
8function avoidFlood(rains: number[]): number[] {
9    const n = rains.length;
10    const result: number[] = new Array(n).fill(-1);
11  
12    // Set of sunny days (when we can dry a lake)
13    const sunnyDays: TreeSet<number> = new TreeSet<number>();
14  
15    // Map tracking the last day each lake was filled
16    const lastRainDay: Map<number, number> = new Map<number, number>();
17  
18    for (let i = 0; i < n; ++i) {
19        const lake = rains[i];
20      
21        if (lake > 0) {
22            // It's raining on this lake
23            if (lastRainDay.has(lake)) {
24                // This lake was already filled, need to dry it before today
25                const previousRainDay = lastRainDay.get(lake)!;
26              
27                // Find the first sunny day after the previous rain
28                const dryingDay = sunnyDays.higher(previousRainDay);
29              
30                if (dryingDay === undefined) {
31                    // No sunny day available to dry this lake - flooding occurs
32                    return [];
33                }
34              
35                // Use this sunny day to dry the lake
36                result[dryingDay] = lake;
37                sunnyDays.delete(dryingDay);
38            }
39          
40            // Update the last rain day for this lake
41            lastRainDay.set(lake, i);
42        } else {
43            // It's a sunny day, we can potentially dry a lake
44            sunnyDays.add(i);
45            result[i] = 1; // Default: dry lake 1 (will be overwritten if needed)
46        }
47    }
48  
49    return result;
50}
51
52// Type definition for comparison function
53type Compare<T> = (lhs: T, rhs: T) => number;
54
55/**
56 * Red-Black Tree Node
57 * Represents a node in the Red-Black Tree with color property for balancing
58 */
59class RBTreeNode<T = number> {
60    data: T;
61    count: number; // Number of duplicates
62    left: RBTreeNode<T> | null;
63    right: RBTreeNode<T> | null;
64    parent: RBTreeNode<T> | null;
65    color: number; // 0 = RED, 1 = BLACK
66  
67    constructor(data: T) {
68        this.data = data;
69        this.left = this.right = this.parent = null;
70        this.color = 0; // New nodes are RED by default
71        this.count = 1;
72    }
73  
74    /**
75     * Get the sibling of this node
76     */
77    sibling(): RBTreeNode<T> | null {
78        if (!this.parent) return null;
79        return this.isOnLeft() ? this.parent.right : this.parent.left;
80    }
81  
82    /**
83     * Check if this node is the left child of its parent
84     */
85    isOnLeft(): boolean {
86        return this === this.parent!.left;
87    }
88  
89    /**
90     * Check if this node has at least one red child
91     */
92    hasRedChild(): boolean {
93        return (
94            Boolean(this.left && this.left.color === 0) ||
95            Boolean(this.right && this.right.color === 0)
96        );
97    }
98}
99
100/**
101 * Red-Black Tree implementation
102 * Self-balancing binary search tree with O(log n) operations
103 */
104class RBTree<T> {
105    root: RBTreeNode<T> | null;
106    lt: (l: T, r: T) => boolean; // Less than comparator
107  
108    constructor(compare: Compare<T> = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0)) {
109        this.root = null;
110        this.lt = (l: T, r: T) => compare(l, r) < 0;
111    }
112  
113    /**
114     * Perform left rotation on the given node
115     */
116    rotateLeft(pivotNode: RBTreeNode<T>): void {
117        const rightChild = pivotNode.right!;
118        pivotNode.right = rightChild.left;
119      
120        if (pivotNode.right) {
121            pivotNode.right.parent = pivotNode;
122        }
123      
124        rightChild.parent = pivotNode.parent;
125      
126        if (!pivotNode.parent) {
127            this.root = rightChild;
128        } else if (pivotNode === pivotNode.parent.left) {
129            pivotNode.parent.left = rightChild;
130        } else {
131            pivotNode.parent.right = rightChild;
132        }
133      
134        rightChild.left = pivotNode;
135        pivotNode.parent = rightChild;
136    }
137  
138    /**
139     * Perform right rotation on the given node
140     */
141    rotateRight(pivotNode: RBTreeNode<T>): void {
142        const leftChild = pivotNode.left!;
143        pivotNode.left = leftChild.right;
144      
145        if (pivotNode.left) {
146            pivotNode.left.parent = pivotNode;
147        }
148      
149        leftChild.parent = pivotNode.parent;
150      
151        if (!pivotNode.parent) {
152            this.root = leftChild;
153        } else if (pivotNode === pivotNode.parent.left) {
154            pivotNode.parent.left = leftChild;
155        } else {
156            pivotNode.parent.right = leftChild;
157        }
158      
159        leftChild.right = pivotNode;
160        pivotNode.parent = leftChild;
161    }
162  
163    /**
164     * Swap colors of two nodes
165     */
166    swapColor(node1: RBTreeNode<T>, node2: RBTreeNode<T>): void {
167        const tempColor = node1.color;
168        node1.color = node2.color;
169        node2.color = tempColor;
170    }
171  
172    /**
173     * Swap data of two nodes
174     */
175    swapData(node1: RBTreeNode<T>, node2: RBTreeNode<T>): void {
176        const tempData = node1.data;
177        node1.data = node2.data;
178        node2.data = tempData;
179    }
180  
181    /**
182     * Fix Red-Black Tree properties after insertion
183     */
184    fixAfterInsert(currentNode: RBTreeNode<T>): void {
185        let parent = null;
186        let grandParent = null;
187      
188        // Fix violations while current node is not root, not black, and parent is red
189        while (currentNode !== this.root && currentNode.color !== 1 && currentNode.parent?.color === 0) {
190            parent = currentNode.parent;
191            grandParent = currentNode.parent.parent;
192          
193            // Case A: Parent is left child of grandparent
194            if (parent === grandParent?.left) {
195                const uncle = grandParent.right;
196              
197                // Case 1: Uncle is red - only recoloring needed
198                if (uncle && uncle.color === 0) {
199                    grandParent.color = 0;
200                    parent.color = 1;
201                    uncle.color = 1;
202                    currentNode = grandParent;
203                } else {
204                    // Case 2: Current node is right child - left rotation needed
205                    if (currentNode === parent.right) {
206                        this.rotateLeft(parent);
207                        currentNode = parent;
208                        parent = currentNode.parent;
209                    }
210                  
211                    // Case 3: Current node is left child - right rotation needed
212                    this.rotateRight(grandParent);
213                    this.swapColor(parent!, grandParent);
214                    currentNode = parent!;
215                }
216            } else {
217                // Case B: Parent is right child of grandparent
218                const uncle = grandParent!.left;
219              
220                // Case 1: Uncle is red - only recoloring needed
221                if (uncle != null && uncle.color === 0) {
222                    grandParent!.color = 0;
223                    parent.color = 1;
224                    uncle.color = 1;
225                    currentNode = grandParent!;
226                } else {
227                    // Case 2: Current node is left child - right rotation needed
228                    if (currentNode === parent.left) {
229                        this.rotateRight(parent);
230                        currentNode = parent;
231                        parent = currentNode.parent;
232                    }
233                  
234                    // Case 3: Current node is right child - left rotation needed
235                    this.rotateLeft(grandParent!);
236                    this.swapColor(parent!, grandParent!);
237                    currentNode = parent!;
238                }
239            }
240        }
241      
242        // Root must always be black
243        this.root!.color = 1;
244    }
245  
246    /**
247     * Delete one instance of the value
248     */
249    delete(value: T): boolean {
250        const node = this.find(value);
251        if (!node) return false;
252      
253        node.count--;
254        if (!node.count) {
255            this.deleteNode(node);
256        }
257        return true;
258    }
259  
260    /**
261     * Delete all instances of the value
262     */
263    deleteAll(value: T): boolean {
264        const node = this.find(value);
265        if (!node) return false;
266      
267        this.deleteNode(node);
268        return true;
269    }
270  
271    /**
272     * Delete a node from the tree
273     */
274    deleteNode(nodeToDelete: RBTreeNode<T>): void {
275        const replacement = findBSTReplacement(nodeToDelete);
276      
277        // Check if both nodes are black
278        const bothBlack = (replacement === null || replacement.color === 1) && nodeToDelete.color === 1;
279        const parent = nodeToDelete.parent!;
280      
281        if (!replacement) {
282            // Node is a leaf
283            if (nodeToDelete === this.root) {
284                this.root = null;
285            } else {
286                if (bothBlack) {
287                    // Both black - fix double black
288                    this.fixDoubleBlack(nodeToDelete);
289                } else {
290                    // One is red - make sibling red if exists
291                    if (nodeToDelete.sibling()) {
292                        nodeToDelete.sibling()!.color = 0;
293                    }
294                }
295              
296                // Remove node from tree
297                if (nodeToDelete.isOnLeft()) {
298                    parent.left = null;
299                } else {
300                    parent.right = null;
301                }
302            }
303            return;
304        }
305      
306        if (!nodeToDelete.left || !nodeToDelete.right) {
307            // Node has one child
308            if (nodeToDelete === this.root) {
309                // Replace root's data and remove child
310                nodeToDelete.data = replacement.data;
311                nodeToDelete.left = nodeToDelete.right = null;
312            } else {
313                // Replace node with its child
314                if (nodeToDelete.isOnLeft()) {
315                    parent.left = replacement;
316                } else {
317                    parent.right = replacement;
318                }
319                replacement.parent = parent;
320              
321                if (bothBlack) {
322                    this.fixDoubleBlack(replacement);
323                } else {
324                    replacement.color = 1;
325                }
326            }
327            return;
328        }
329      
330        // Node has two children - swap with successor and recurse
331        this.swapData(replacement, nodeToDelete);
332        this.deleteNode(replacement);
333      
334        /**
335         * Find replacement node in BST deletion
336         */
337        function findBSTReplacement(node: RBTreeNode<T>): RBTreeNode<T> | null {
338            // Two children: find inorder successor
339            if (node.left && node.right) {
340                return findSuccessor(node.right);
341            }
342          
343            // Leaf node
344            if (!node.left && !node.right) {
345                return null;
346            }
347          
348            // Single child
349            return node.left ?? node.right;
350        }
351      
352        /**
353         * Find inorder successor (leftmost node in right subtree)
354         */
355        function findSuccessor(node: RBTreeNode<T>): RBTreeNode<T> {
356            let current = node;
357            while (current.left) {
358                current = current.left;
359            }
360            return current;
361        }
362    }
363  
364    /**
365     * Fix double black violation in Red-Black Tree
366     */
367    fixDoubleBlack(node: RBTreeNode<T>): void {
368        // Base case: reached root
369        if (node === this.root) return;
370      
371        const sibling = node.sibling();
372        const parent = node.parent!;
373      
374        if (!sibling) {
375            // No sibling - push double black up
376            this.fixDoubleBlack(parent);
377        } else {
378            if (sibling.color === 0) {
379                // Red sibling
380                parent.color = 0;
381                sibling.color = 1;
382              
383                if (sibling.isOnLeft()) {
384                    this.rotateRight(parent);
385                } else {
386                    this.rotateLeft(parent);
387                }
388              
389                this.fixDoubleBlack(node);
390            } else {
391                // Black sibling
392                if (sibling.hasRedChild()) {
393                    // At least one red child
394                    if (sibling.left && sibling.left.color === 0) {
395                        if (sibling.isOnLeft()) {
396                            // Left-left case
397                            sibling.left.color = sibling.color;
398                            sibling.color = parent.color;
399                            this.rotateRight(parent);
400                        } else {
401                            // Right-left case
402                            sibling.left.color = parent.color;
403                            this.rotateRight(sibling);
404                            this.rotateLeft(parent);
405                        }
406                    } else {
407                        if (sibling.isOnLeft()) {
408                            // Left-right case
409                            sibling.right!.color = parent.color;
410                            this.rotateLeft(sibling);
411                            this.rotateRight(parent);
412                        } else {
413                            // Right-right case
414                            sibling.right!.color = sibling.color;
415                            sibling.color = parent.color;
416                            this.rotateLeft(parent);
417                        }
418                    }
419                    parent.color = 1;
420                } else {
421                    // Two black children
422                    sibling.color = 0;
423                  
424                    if (parent.color === 1) {
425                        this.fixDoubleBlack(parent);
426                    } else {
427                        parent.color = 1;
428                    }
429                }
430            }
431        }
432    }
433  
434    /**
435     * Insert a value into the tree
436     */
437    insert(data: T): boolean {
438        // Find insertion position
439        let parent = this.root;
440        while (parent) {
441            if (this.lt(data, parent.data)) {
442                if (!parent.left) break;
443                else parent = parent.left;
444            } else if (this.lt(parent.data, data)) {
445                if (!parent.right) break;
446                else parent = parent.right;
447            } else break;
448        }
449      
450        // Create and insert new node
451        const newNode = new RBTreeNode(data);
452      
453        if (!parent) {
454            this.root = newNode;
455        } else if (this.lt(newNode.data, parent.data)) {
456            parent.left = newNode;
457        } else if (this.lt(parent.data, newNode.data)) {
458            parent.right = newNode;
459        } else {
460            // Duplicate value - increment count
461            parent.count++;
462            return false;
463        }
464      
465        newNode.parent = parent;
466        this.fixAfterInsert(newNode);
467        return true;
468    }
469  
470    /**
471     * Find a node with the given value
472     */
473    find(data: T): RBTreeNode<T> | null {
474        let current = this.root;
475      
476        while (current) {
477            if (this.lt(data, current.data)) {
478                current = current.left;
479            } else if (this.lt(current.data, data)) {
480                current = current.right;
481            } else {
482                break;
483            }
484        }
485      
486        return current ?? null;
487    }
488  
489    /**
490     * Generator for in-order traversal
491     */
492    *inOrder(root: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
493        if (!root) return;
494      
495        yield* this.inOrder(root.left!);
496        yield root.data;
497        yield* this.inOrder(root.right!);
498    }
499  
500    /**
501     * Generator for reverse in-order traversal
502     */
503    *reverseInOrder(root: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
504        if (!root) return;
505      
506        yield* this.reverseInOrder(root.right!);
507        yield root.data;
508        yield* this.reverseInOrder(root.left!);
509    }
510}
511
512/**
513 * TreeSet implementation using Red-Black Tree
514 * Provides ordered set operations with O(log n) complexity
515 */
516class TreeSet<T = number> {
517    private _size: number;
518    private tree: RBTree<T>;
519    private compare: Compare<T>;
520  
521    constructor(
522        collection: T[] | Compare<T> = [],
523        compare: Compare<T> = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0),
524    ) {
525        // Handle overloaded constructor
526        if (typeof collection === 'function') {
527            compare = collection;
528            collection = [];
529        }
530      
531        this._size = 0;
532        this.compare = compare;
533        this.tree = new RBTree(compare);
534      
535        // Initialize with collection
536        for (const value of collection) {
537            this.add(value);
538        }
539    }
540  
541    /**
542     * Get the size of the set
543     */
544    size(): number {
545        return this._size;
546    }
547  
548    /**
549     * Check if value exists in the set
550     */
551    has(value: T): boolean {
552        return !!this.tree.find(value);
553    }
554  
555    /**
556     * Add a value to the set
557     */
558    add(value: T): boolean {
559        const wasAdded = this.tree.insert(value);
560        this._size += wasAdded ? 1 : 0;
561        return wasAdded;
562    }
563  
564    /**
565     * Delete a value from the set
566     */
567    delete(value: T): boolean {
568        const wasDeleted = this.tree.deleteAll(value);
569        this._size -= wasDeleted ? 1 : 0;
570        return wasDeleted;
571    }
572  
573    /**
574     * Find the smallest value >= the given value
575     */
576    ceil(value: T): T | undefined {
577        let current = this.tree.root;
578        let ceiling = null;
579      
580        while (current) {
581            if (this.compare(current.data, value) >= 0) {
582                ceiling = current;
583                current = current.left;
584            } else {
585                current = current.right;
586            }
587        }
588      
589        return ceiling?.data;
590    }
591  
592    /**
593     * Find the largest value <= the given value
594     */
595    floor(value: T): T | undefined {
596        let current = this.tree.root;
597        let floor = null;
598      
599        while (current) {
600            if (this.compare(value, current.data) >= 0) {
601                floor = current;
602                current = current.right;
603            } else {
604                current = current.left;
605            }
606        }
607      
608        return floor?.data;
609    }
610  
611    /**
612     * Find the smallest value > the given value
613     */
614    higher(value: T): T | undefined {
615        let current = this.tree.root;
616        let higher = null;
617      
618        while (current) {
619            if (this.compare(value, current.data) < 0) {
620                higher = current;
621                current = current.left;
622            } else {
623                current = current.right;
624            }
625        }
626      
627        return higher?.data;
628    }
629  
630    /**
631     * Find the largest value < the given value
632     */
633    lower(value: T): T | undefined {
634        let current = this.tree.root;
635        let lower = null;
636      
637        while (current) {
638            if (this.compare(current.data, value) < 0) {
639                lower = current;
640                current = current.right;
641            } else {
642                current = current.left;
643            }
644        }
645      
646        return lower?.data;
647    }
648  
649    /**
650     * Get the minimum value in the set
651     */
652    first(): T | undefined {
653        return this.tree.inOrder().next().value;
654    }
655  
656    /**
657     * Get the maximum value in the set
658     */
659    last(): T | undefined {
660        return this.tree.reverseInOrder().next().value;
661    }
662  
663    /**
664     * Remove and return the minimum value
665     */
666    shift(): T | undefined {
667        const firstValue = this.first();
668        if (firstValue === undefined) return undefined;
669      
670        this.delete(firstValue);
671        return firstValue;
672    }
673  
674    /**
675     * Remove and return the maximum value
676     */
677    pop(): T | undefined {
678        const lastValue = this.last();
679        if (lastValue === undefined) return undefined;
680      
681        this.delete(lastValue);
682        return lastValue;
683    }
684  
685    /**
686     * Iterator for the set
687     */
688    *[Symbol.iterator](): Generator<T, void, void> {
689        yield* this.values();
690    }
691  
692    /**
693     * Generator for keys (same as values for a set)
694     */
695    *keys(): Generator<T, void, void> {
696        yield* this.values();
697    }
698  
699    /**
700     * Generator for values in ascending order
701     */
702    *values(): Generator<T, undefined, void> {
703        yield* this.tree.inOrder();
704        return undefined;
705    }
706  
707    /**
708     * Generator for values in descending order
709     */
710    *rvalues(): Generator<T, undefined, void> {
711        yield* this.tree.reverseInOrder();
712        return undefined;
713    }
714}
715

Time and Space Complexity

Time Complexity: O(n × log n)

The algorithm iterates through the rains array once with n elements. For each element:

When v != 0 (rainy day):
- Checking if v in rainy takes O(1) with a hash map
- bisect_right operation on the SortedList takes O(log n)
- discard operation on the SortedList takes O(log n)
- Updating the hash map takes O(1)
When v == 0 (sunny day):
- add operation on the SortedList takes O(log n)

Since we perform O(log n) operations for each of the n elements in the worst case, the overall time complexity is O(n × log n).

Space Complexity: O(n)

The algorithm uses:

ans array: O(n) space
sunny SortedList: can contain at most n elements, so O(n) space
rainy dictionary: can contain at most n unique lakes, so O(n) space

The total space complexity is O(n) + O(n) + O(n) = O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Using Regular List Instead of SortedList

Pitfall: Using a regular Python list for sunny_days and manually sorting it each time or using bisect on an unsorted list.

# WRONG: This won't maintain sorted order
sunny_days = []
sunny_days.append(day)  # List may become unsorted
idx = bisect.bisect_right(sunny_days, last_rain_day[lake])  # Incorrect on unsorted list

Why it fails: Binary search requires a sorted collection. If you use a regular list and forget to sort it after each insertion, bisect_right will give incorrect results.

Solution: Use SortedList from sortedcontainers or maintain sorting manually:

# Option 1: Use SortedList (recommended)
from sortedcontainers import SortedList
sunny_days = SortedList()

# Option 2: Use set and convert to sorted list when needed
sunny_days = set()
# When searching:
sorted_sunny = sorted(sunny_days)
idx = bisect.bisect_right(sorted_sunny, last_rain_day[lake])

2. Incorrect Removal from Collection

Pitfall: Using wrong method to remove elements or accessing by wrong index.

# WRONG: Trying to remove by index instead of value
sunny_days.pop(sunny_day_idx)  # This removes at index position

# WRONG: Not checking if element exists before removal
sunny_days.remove(dry_day)  # Raises error if not found

Solution: Use the correct removal method:

# CORRECT: Remove by value, not index
dry_day = sunny_days[sunny_day_idx]
sunny_days.discard(dry_day)  # or sunny_days.remove(dry_day)

3. Off-by-One Error in Binary Search

Pitfall: Using bisect_left instead of bisect_right.

# WRONG: bisect_left might return the exact rain day
sunny_day_idx = sunny_days.bisect_left(last_rain_day[lake])

Why it fails: If there's a sunny day on the same day as the last rain (which shouldn't happen in valid input), bisect_left would return that index. We need the first sunny day after the last rain, so bisect_right is correct.

Solution: Always use bisect_right to find the first position greater than the target:

# CORRECT: Gets first sunny day strictly after last rain
sunny_day_idx = sunny_days.bisect_right(last_rain_day[lake])

4. Not Handling Edge Cases for Result Array

Pitfall: Forgetting to set a valid lake number for unused sunny days.

# WRONG: Leaving result[day] as -1 for sunny days
if lake == 0:
    sunny_days.add(day)
    # Forgot to set result[day] to a valid lake number

Why it fails: The problem requires that on sunny days, you must specify which lake to dry (even if it doesn't matter). Returning -1 for a sunny day is invalid.

Solution: Always set a valid lake number (typically 1) for sunny days:

# CORRECT: Set a default lake to dry
if lake == 0:
    sunny_days.add(day)
    result[day] = 1  # Can be any valid lake number

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which type of traversal does breadth first search do?

Pre-order traversal

Post-order traversal

In-order traversal

Level-order traversal

I don't know