LeetCode 3840 - House Robber V

LeetCode Problem 3840

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming

Solution

Problem Understanding

The problem is a variation of the classic House Robber problem, but with an additional constraint: each house has a color, and you cannot rob two adjacent houses if they share the same color. We are given two arrays nums and colors of length n. The array nums represents the amount of money at each house, while colors represents the color code of each house. The goal is to maximize the total amount of money robbed under the constraint that two adjacent houses with the same color cannot both be robbed.

The expected output is a single integer representing the maximum money that can be robbed while respecting the color-adjacency constraint. The constraints indicate that the input can be large (n up to 10^5), so any solution that is quadratic or worse in n will be too slow. This implies the need for a linear-time dynamic programming solution.

Important edge cases include:

• All houses having the same color, which limits robbing to non-adjacent houses only.
• Alternating colors, which may allow robbing every house.
• A single house, where the answer is simply the amount in that house.
• Large nums values up to 10^5, which may require handling of large sums (integer overflow is only relevant in languages like Go).

Approaches

The brute-force approach is to recursively consider each house with the choice of robbing or skipping it, while checking the color constraint against the previous house. This gives the correct answer because it explores all valid combinations, but it has exponential time complexity, O(2^n), which is impractical for n up to 10^5.

The key insight for an optimal solution is dynamic programming. We maintain the maximum money that can be robbed ending at each house. To avoid violating the color-adjacency constraint, we track the maximum money robbed up to the previous house for each color. At each house, we consider whether we can rob it based on the previous house’s color and update a mapping from color to maximum robbed money. This approach works in linear time because we process each house once and perform constant-time updates using a hash map.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(2^n)	O(n)	Recursively explore all subsets, checking color constraints
Optimal	O(n)	O(c)	Dynamic programming with hash map to track color-specific max, c is number of distinct colors

Algorithm Walkthrough

1. Initialize a hash map last_color_max to track the maximum money robbed up to the previous house for each color. Also initialize a variable overall_max to track the global maximum robbed money.
2. Iterate through each house i from 0 to n-1.
3. For house i, consider robbing it. If the previous house has the same color, we cannot include it in a consecutive sum with the same color. Therefore, calculate current_max as the maximum of overall_max minus the previous maximum of this color, plus nums[i].
4. Update last_color_max[colors[i]] to the maximum of its current value and current_max. This ensures we always keep the best value ending at this color.
5. Update overall_max to the maximum of itself and current_max.
6. After processing all houses, overall_max contains the maximum money that can be robbed under the constraints.

Why it works: The algorithm maintains two invariants. First, last_color_max[color] always contains the maximum money that can be robbed ending at a house of that color. Second, overall_max contains the maximum money robbed across all houses. By updating these at each step, the algorithm guarantees that all valid sequences are considered, and the maximum is correctly computed.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def rob(self, nums: List[int], colors: List[int]) -> int:
        last_color_max = defaultdict(int)
        overall_max = 0
        
        for i in range(len(nums)):
            color = colors[i]
            current_max = overall_max
            if color in last_color_max:
                current_max -= last_color_max[color]
            current_max += nums[i]
            last_color_max[color] = max(last_color_max[color], current_max)
            overall_max = max(overall_max, current_max)
        
        return overall_max

This implementation uses a hash map last_color_max to track the maximum money robbed for each color, ensuring that consecutive houses with the same color do not violate the constraint. The variable overall_max keeps track of the global maximum, and we update it iteratively. Each house is processed once, making this solution linear in time.

Go Solution

func rob(nums []int, colors []int) int64 {
    lastColorMax := make(map[int]int64)
    var overallMax int64 = 0
    
    for i := 0; i < len(nums); i++ {
        color := colors[i]
        currentMax := overallMax
        if val, exists := lastColorMax[color]; exists {
            currentMax -= val
        }
        currentMax += int64(nums[i])
        if val, exists := lastColorMax[color]; !exists || currentMax > val {
            lastColorMax[color] = currentMax
        }
        if currentMax > overallMax {
            overallMax = currentMax
        }
    }
    
    return overallMax
}

In Go, we use a map[int]int64 to store the color-specific maximums. We cast nums[i] to int64 to handle potential integer overflow. Otherwise, the logic mirrors the Python version.

Worked Examples

Example 1: nums = [1,4,3,5], colors = [1,1,2,2]

i	nums[i]	colors[i]	last_color_max	overall_max	current_max
0	1	1	{}	0	1
1	4	1	{1:1}	1	4
2	3	2	{1:4}	4	7
3	5	2	{1:4, 2:7}	7	9

Result: 9

Example 2: nums = [3,1,2,4], colors = [2,3,2,2]

i	nums[i]	colors[i]	last_color_max	overall_max	current_max
0	3	2	{}	0	3
1	1	3	{2:3}	3	4
2	2	2	{2:3,3:4}	4	5
3	4	2	{2:5,3:4}	5	8

Result: 8

Example 3: nums = [10,1,3,9], colors = [1,1,1,2]

i	nums[i]	colors[i]	last_color_max	overall_max	current_max
0	10	1	{}	0	10
1	1	1	{1:10}	10	1
2	3	1	{1:10}	10	13
3	9	2	{1:13,2:0}	13	22

Result: 22

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each house is processed exactly once; hash map operations are amortized O(1)
Space	O(c)	We maintain a hash map of size equal to the number of distinct colors c

The linear time complexity is feasible for n up to 10^5, and space usage is reasonable since the number of colors is bounded by n.

Test Cases

# provided examples
assert Solution().rob([1,4,3,5], [1,1,2,2]) == 9
assert Solution().rob([3,1,2,4], [2,3,2,2]) == 8
assert Solution().rob([10,1,3,9], [1,1,1,2]) == 22

# edge cases
assert Solution().rob([5], [1]) == 5  # single house
assert Solution().rob([1,2,3,4,5], [1,1,1,1,1]) == 9  # same color, pick non-adjacent
assert Solution().rob([1,2,3,4,5], [1,2,1,2,1]) == 15  # alternating colors, pick all
assert Solution().rob([100000,100000], [1,1]) == 100000  # large nums with same color

Test	Why
[1,4,3,5], [1,1,2,2]	Standard example with adjacent same-color restriction
[3,1,