Problem Description

You are given a 0-indexed integer array nums. Your task is to determine if you can partition this array into one or more contiguous subarrays where each subarray must satisfy exactly one of the following conditions:

1. Two equal elements: The subarray contains exactly 2 elements that are the same. For example, [2,2] is valid.

2. Three equal elements: The subarray contains exactly 3 elements that are all the same. For example, [4,4,4] is valid.

3. Three consecutive increasing elements: The subarray contains exactly 3 elements where each element is exactly 1 more than the previous element. For example, [3,4,5] is valid (since 4-3=1 and 5-4=1), but [1,3,5] is not valid.

The partition must use all elements in the array, and the subarrays must be contiguous (you cannot skip elements or rearrange them).

Return true if at least one valid partition exists for the entire array, otherwise return false.

For example:

• nums = [4,4,4,5,6] can be partitioned as [4,4,4] and [5,6] (but [5,6] doesn't satisfy any condition, so this would return false)
• nums = [1,1,1,2] can be partitioned as [1,1,1] and [2] (but a single element doesn't satisfy any condition, so this would return false)
• nums = [1,1,2,2,3,3] can be partitioned as [1,1], [2,2], [3,3] (all subarrays satisfy condition 1, so this would return true)

Intuition

The key insight is that we need to make a series of decisions about how to partition the array. At each position, we need to decide: should we form a 2-element subarray, a 3-element subarray, or is it even possible to form a valid subarray starting from this position?

Think of it like walking through the array from left to right. At each step, we have limited choices - we can only "consume" the next 2 or 3 elements if they form a valid pattern. Once we make a choice, we move forward and face the same decision at the new position.

This naturally leads to a recursive approach: at position i, we check what valid subarrays we can form starting from i, and for each valid option, we recursively check if the remaining array (starting from where our current subarray ends) can also be validly partitioned.

The problem exhibits optimal substructure - if we can validly partition the array from position i onwards, it doesn't matter how we got to position i. This means we'll likely encounter the same subproblem multiple times. For instance, different partitioning choices in the beginning might lead us to check "can we partition from index 5 onwards?" multiple times.

This overlap in subproblems suggests using dynamic programming or memoization. We can cache the result for each starting position i - once we know whether the array from position i onwards can be validly partitioned, we don't need to recalculate it.

The base case is straightforward: if we've successfully consumed all elements (reached or passed the end of the array), we've found a valid partition. If we're stuck at a position where no valid subarray can be formed, that path fails. The answer is whether starting from position 0, we can find at least one way to validly partition the entire array.

Learn more about Dynamic Programming patterns.

Solution Approach

We implement a recursive function dfs(i) that determines whether a valid partition exists starting from index i. The function uses memoization with @cache decorator to avoid recalculating the same subproblems.

Base Case:

• If i >= n (we've reached or passed the end of the array), return True as we've successfully partitioned all elements.

Recursive Cases: At each position i, we check three possible valid subarray patterns:

1. Two equal elements (a): Check if i+1 < n and nums[i] == nums[i+1]. If true, we can form a 2-element subarray and recursively check dfs(i+2).

2. Three equal elements (b): Check if i+2 < n and nums[i] == nums[i+1] == nums[i+2]. If true, we can form a 3-element subarray and recursively check dfs(i+3).

3. Three consecutive increasing elements (c): Check if i+2 < n, nums[i+1] - nums[i] == 1, and nums[i+2] - nums[i+1] == 1. If true, we can form a 3-element consecutive subarray and recursively check dfs(i+3).

Decision Logic: The function returns True if:

• Pattern a is valid AND dfs(i+2) returns True, OR
• Either pattern b or c is valid AND dfs(i+3) returns True

This can be expressed as:

dfs(i) = (a and dfs(i+2)) or ((b or c) and dfs(i+3))

Implementation Details:

• We store the array length n for boundary checking
• The @cache decorator automatically memoizes results, preventing redundant calculations
• We start the recursion with dfs(0) to check if the entire array can be validly partitioned

Time Complexity: O(n) where n is the length of the array, since each position is computed at most once due to memoization.

Space Complexity: O(n) for the recursion stack and memoization cache.

Example Walkthrough

Let's walk through the example nums = [1,1,2,3,4] to illustrate how the solution works.

Initial Setup:

• Array: [1,1,2,3,4], length n = 5
• We start with dfs(0) to check if we can partition from index 0

Step 1: dfs(0) - Processing index 0 (value = 1)

• Check pattern a (two equal): nums[0] == nums[1]? → 1 == 1 ✓ Valid
• Check pattern b (three equal): nums[0] == nums[1] == nums[2]? → 1 == 1 == 2 ✗ Invalid
• Check pattern c (three consecutive): nums[1]-nums[0] == 1 and nums[2]-nums[1] == 1? → 1-1 == 1 ✗ Invalid
• Since pattern a is valid, we check: dfs(2)

Step 2: dfs(2) - Processing index 2 (value = 2)

• Check pattern a (two equal): nums[2] == nums[3]? → 2 == 3 ✗ Invalid
• Check pattern b (three equal): nums[2] == nums[3] == nums[4]? → 2 == 3 == 4 ✗ Invalid
• Check pattern c (three consecutive): nums[3]-nums[2] == 1 and nums[4]-nums[3] == 1? → 3-2 == 1 ✓ and 4-3 == 1 ✓ Valid
• Since pattern c is valid, we check: dfs(5)

Step 3: dfs(5) - Processing index 5

• Since 5 >= n (5 >= 5), we've successfully consumed all elements
• Return True

Backtracking the results:

• dfs(5) returns True
• dfs(2) returns True (because pattern c was valid and dfs(5) returned True)
• dfs(0) returns True (because pattern a was valid and dfs(2) returned True)

Final Result: True

The array [1,1,2,3,4] can be validly partitioned as:

• [1,1] - satisfies condition 1 (two equal elements)
• [2,3,4] - satisfies condition 3 (three consecutive increasing elements)

The memoization ensures that if we had encountered the same index multiple times through different paths, we would only compute it once and reuse the cached result.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm uses dynamic programming with memoization through the @cache decorator. Each index i from 0 to n-1 is visited at most once due to memoization. When we reach an index i, we perform constant time operations:

• Checking boundary conditions (i + 1 < n, i + 2 < n)
• Comparing array elements (nums[i] == nums[i + 1], etc.)
• Computing differences (nums[i + 1] - nums[i])

After these O(1) checks, we make recursive calls to either dfs(i + 2) or dfs(i + 3). Since each state is cached, every index is processed exactly once, giving us a total of n states to compute. Therefore, the overall time complexity is O(n).

Space Complexity: O(n)

The space complexity consists of two components:

1. Recursion Stack: In the worst case, the recursion depth can go up to O(n) when we always take 2-element partitions (calling dfs(i + 2) repeatedly), resulting in approximately n/2 recursive calls.
2. Memoization Cache: The @cache decorator stores the result for each unique value of i. Since i ranges from 0 to n-1, we store at most n different states in the cache.

Both the recursion stack and the cache contribute O(n) space, so the total space complexity is O(n).

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

Common Pitfalls

1. Incorrect Condition Checking Order

A common mistake is checking conditions without proper boundary validation first, leading to index out of bounds errors.

Incorrect approach:

# Wrong - accesses nums[i+1] before checking if i+1 < n
has_two_equal = nums[i] == nums[i+1] and i + 1 < n

Correct approach:

# Right - checks boundary first with short-circuit evaluation
has_two_equal = i + 1 < n and nums[i] == nums[i+1]

2. Greedy Algorithm Temptation

Many developers initially try a greedy approach, always choosing the longest valid subarray at each position. This fails because sometimes taking a shorter valid subarray leads to the overall solution.

Example where greedy fails:

nums = [4,4,4,5,6,7]
# Greedy would take [4,4,4] first, leaving [5,6,7] which is valid
# But what if nums = [4,4,5,6,7,7]?
# Greedy takes [4,4] then [5,6,7], leaving single [7] - invalid!
# Correct partition: [4,4], [5,6,7], [7] - wait, that's still invalid
# Actually correct: Can't be partitioned, but point is greedy doesn't work

3. Missing Memoization

Without memoization, the recursive solution has exponential time complexity due to overlapping subproblems.

Without memoization (TLE for large inputs):

def dfs(i):
    if i >= n:
        return True
    # ... check conditions ...
    return result  # Same states computed multiple times

With memoization (efficient):

@cache
def dfs(i):
    # Same logic, but results are cached

4. Incorrect Consecutive Check

A subtle bug is checking consecutive elements incorrectly, such as only checking the difference between first and last elements.

Incorrect:

# Wrong - doesn't ensure ALL differences are 1
has_consecutive = i + 2 < n and nums[i+2] - nums[i] == 2

Correct:

# Right - checks each adjacent pair
has_consecutive = (i + 2 < n and 
                  nums[i+1] - nums[i] == 1 and 
                  nums[i+2] - nums[i+1] == 1)

5. Off-by-One Errors in Base Case

The base case should return True when i >= n, not just when i == n, because after taking the last valid subarray, the index might jump past the array length.

Incorrect:

if i == n:  # Misses case when last subarray is 3 elements
    return True

Correct:

if i >= n:  # Handles all cases properly
    return True