LeetCode 2261 - K Divisible Elements Subarrays

This problem requires counting the number of distinct subarrays of a given integer array nums that satisfy a specific constraint: each subarray can contain at most k elements divisible by p.

leetcodemediumarrayhash-tabletrierolling-hashhash-functionenumeration

LeetCode Problem 2261

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Trie, Rolling Hash, Hash Function, Enumeration

Solution

Problem Understanding

This problem requires counting the number of distinct subarrays of a given integer array nums that satisfy a specific constraint: each subarray can contain at most k elements divisible by p. A subarray is a contiguous slice of the array, and "distinct" means either differing in length or differing in at least one element.

For example, consider nums = [2,3,3,2,2], k = 2, and p = 2. Here, elements divisible by p are 2 (indices 0, 3, 4). Any subarray containing more than 2 elements divisible by 2 is invalid. The task is to enumerate all unique subarrays meeting this criterion. Even if [2,3] appears multiple times in the array, it counts only once.

The constraints tell us that the array is small (1 <= nums.length <= 200), meaning an $O(n^2)$ solution is feasible, but brute-force approaches may still need careful optimization to avoid counting duplicates inefficiently. Key edge cases include arrays where all elements are divisible by p, arrays with no divisible elements, arrays with repeated numbers, and situations where k equals n or is very small (like 1).

Approaches

The brute-force approach is straightforward: generate all possible subarrays, count the elements divisible by p in each, and store them in a hash set to ensure uniqueness. This is correct but inefficient because the number of subarrays is $O(n^2)$, and comparing arrays for uniqueness adds overhead, leading to $O(n^3)$ time if not optimized.

The optimal approach uses the observation that we can generate subarrays in a sliding-window fashion, counting divisible elements on the fly. We maintain a hash set of tuples representing subarrays to efficiently avoid duplicates. By stopping the extension of a subarray once the divisible count exceeds k, we prune unnecessary computation. This reduces the effective operations to $O(n^2)$ since each subarray is checked at most once, and Python or Go hash sets handle uniqueness efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^3) O(n^2) Generate all subarrays, count divisible, store in set of arrays
Optimal O(n^2) O(n^2) Sliding window + early stop + hash set for uniqueness

Algorithm Walkthrough

  1. Initialize an empty hash set to store unique subarrays.
  2. Iterate over each starting index i of nums.
  3. For each starting index, initialize a counter count_div to track the number of elements divisible by p.
  4. Iterate over each ending index j starting from i to the end of the array.
  5. For each nums[j], increment count_div if nums[j] % p == 0.
  6. If count_div > k, break the inner loop, since any longer subarray from this start will also exceed k.
  7. Otherwise, add the tuple nums[i:j+1] to the hash set to ensure distinct subarrays.
  8. After processing all subarrays, return the size of the hash set, representing the count of distinct valid subarrays.

Why it works: The algorithm generates all valid contiguous subarrays, and the use of a hash set ensures duplicates are ignored. Early termination avoids unnecessary iterations for subarrays that cannot satisfy the condition.

Python Solution

from typing import List

class Solution:
    def countDistinct(self, nums: List[int], k: int, p: int) -> int:
        seen = set()
        n = len(nums)
        for i in range(n):
            count_div = 0
            subarray = []
            for j in range(i, n):
                if nums[j] % p == 0:
                    count_div += 1
                if count_div > k:
                    break
                subarray.append(nums[j])
                seen.add(tuple(subarray))
        return len(seen)

The Python implementation iterates over all starting indices, maintains a count of divisible elements, and builds subarrays incrementally. The tuple(subarray) ensures hashable storage in the set. Breaking early prevents subarrays exceeding the limit from being considered.

Go Solution

func countDistinct(nums []int, k int, p int) int {
    seen := make(map[string]struct{})
    n := len(nums)

    for i := 0; i < n; i++ {
        countDiv := 0
        subarray := []int{}
        for j := i; j < n; j++ {
            if nums[j]%p == 0 {
                countDiv++
            }
            if countDiv > k {
                break
            }
            subarray = append(subarray, nums[j])
            key := fmt.Sprint(subarray)
            seen[key] = struct{}{}
        }
    }
    return len(seen)
}

In Go, we use a map with string keys to emulate a set. fmt.Sprint(subarray) converts the slice to a string for hashable storage. This handles duplicates correctly while maintaining an O(n^2) time complexity.

Worked Examples

Example 1: nums = [2,3,3,2,2], k = 2, p = 2

i j Subarray Divisible count Add to set?
0 0 [2] 1 Yes
0 1 [2,3] 1 Yes
0 2 [2,3,3] 1 Yes
0 3 [2,3,3,2] 2 Yes
0 4 [2,3,3,2,2] 3 No (break)
1 1 [3] 0 Yes
1 2 [3,3] 0 Yes
1 3 [3,3,2] 1 Yes
1 4 [3,3,2,2] 2 Yes
2 2 [3] 0 Already in set
2 3 [3,2] 1 Yes
2 4 [3,2,2] 2 Yes
3 3 [2] 1 Already in set
3 4 [2,2] 2 Yes
4 4 [2] 1 Already in set

Distinct valid subarrays count = 11

Example 2: nums = [1,2,3,4], k = 4, p = 1

All elements divisible, but k >= n. All subarrays included, total = 10.

Complexity Analysis

Measure Complexity Explanation
Time O(n^2) Each subarray starting index is considered, inner loop breaks early when divisible count exceeds k
Space O(n^2) Hash set stores at most O(n^2) distinct subarrays

Because the array length is at most 200, this quadratic approach is acceptable.

Test Cases

# provided examples
assert Solution().countDistinct([2,3,3,2,2], 2, 2) == 11  # Example 1
assert Solution().countDistinct([1,2,3,4], 4, 1) == 10     # Example 2

# edge cases
assert Solution().countDistinct([1,1,1,1], 2, 1) == 10      # repeated elements
assert Solution().countDistinct([2,4,6], 1, 2) == 3        # all divisible, k small
assert Solution().countDistinct([1,3,5,7], 0, 2) == 10     # no divisible elements, k=0
assert Solution().countDistinct([100], 1, 100) == 1        # single element divisible
assert Solution().countDistinct([100], 0, 100) == 0        # single element, but k=0
Test Why
[2,3,3,2,2], 2, 2 Validates the main example with duplicates
[1,2,3,4], 4, 1 All elements divisible, checks k >= n handling
[1,1,1,1], 2, 1 Repeated elements, uniqueness check
[2,4,6], 1, 2 All divisible, early break validation
[1,3,5,7], 0, 2 No divisible elements, k=0 edge case
[100], 1, 100 Single element divisible
[100], 0, 100 Single element but k=0, should return 0

Edge Cases

The first important edge case is when all elements are divisible by p. In this situation, k limits the subarrays, and the early break