K Divisible Elements Subarrays
MedPrompt
Given an integer array nums and two integers k and p, return the number of distinct subarrays, which have at most k elements that are divisible by p.
Two arrays nums1 and nums2 are said to be distinct if:
- They are of different lengths, or
- There exists at least one index
iwherenums1[i] != nums2[i].
A subarray is defined as a non-empty contiguous sequence of elements in an array.
Example 1:
Input: nums = [2,3,3,2,2], k = 2, p = 2
Output: 11
Explanation:
The elements at indices 0, 3, and 4 are divisible by p = 2.
The 11 distinct subarrays which have at most k = 2 elements divisible by 2 are:
[2], [2,3], [2,3,3], [2,3,3,2], [3], [3,3], [3,3,2], [3,3,2,2], [3,2], [3,2,2], and [2,2].
Note that the subarrays [2] and [3] occur more than once in nums, but they should each be counted only once.
The subarray [2,3,3,2,2] should not be counted because it has 3 elements that are divisible by 2.Example 2:
Input: nums = [1,2,3,4], k = 4, p = 1
Output: 10
Explanation:
All element of nums are divisible by p = 1.
Also, every subarray of nums will have at most 4 elements that are divisible by 1.
Since all subarrays are distinct, the total number of subarrays satisfying all the constraints is 10.
Constraints:
1 <= nums.length <= 2001 <= nums[i], p <= 2001 <= k <= nums.length
Follow up:
Can you solve this problem in O(n2) time complexity?
Approaches
3 approaches with complexity analysis and trade-offs.
This approach involves generating every possible subarray of nums, checking if it satisfies the condition (at most k elements divisible by p), and then storing the valid, unique subarrays in a HashSet. The final answer is the size of the set. This is the most straightforward and intuitive way to solve the problem, but it is not the most efficient.
Algorithm
- Initialize an empty
HashSetto store distinct valid subarrays, for example, asSet<List<Integer>>. - Iterate through all possible start indices
iof a subarray from0ton-1. - For each
i, iterate through all possible end indicesjfromiton-1. - For each subarray
nums[i...j], create a temporary list and count the number of elements divisible byp. - A more optimized way is to build the list and count incrementally. For a fixed
i, asjincreases, we appendnums[j]to a running list and update a running count of divisible elements. - If the count of divisible elements is at most
k, add a copy of the current subarray list to theHashSet. The set automatically handles uniqueness. - If the count exceeds
k, we can break the inner loop, as any further extension of the subarray from startiwill also be invalid. - The final answer is the size of the
HashSet.
Walkthrough
We use two nested loops to define the start (i) and end (j) of each subarray. For each subarray nums[i...j], we check its validity. To handle the "distinct" requirement, we add each valid subarray to a HashSet. A HashSet<List<Integer>> works well in Java because List has a content-based equals and hashCode implementation.
The process is as follows: we iterate with a start index i. For each i, we start building a new subarray. A second loop with index j extends this subarray one element at a time. We maintain a count of elements divisible by p for the current subarray nums[i...j]. If this count is within the limit k, we add a copy of the current subarray list to our set. If the count exceeds k, we can stop extending from i because all subsequent subarrays will also be invalid.
import java.util.ArrayList;import java.util.HashSet;import java.util.List;import java.util.Set; class Solution { public int countDistinct(int[] nums, int k, int p) { Set<List<Integer>> distinctSubarrays = new HashSet<>(); int n = nums.length; for (int i = 0; i < n; i++) { int divisibleCount = 0; List<Integer> currentSubarray = new ArrayList<>(); for (int j = i; j < n; j++) { currentSubarray.add(nums[j]); if (nums[j] % p == 0) { divisibleCount++; } if (divisibleCount <= k) { // A new list must be created because the set stores a reference. distinctSubarrays.add(new ArrayList<>(currentSubarray)); } else { // Optimization: if count exceeds k, any further extension is also invalid. break; } } } return distinctSubarrays.size(); }}Complexity
Time
O(n^3). The two nested loops to generate subarrays run in `O(n^2)`. Inside the inner loop, creating a copy of the `currentSubarray` and adding it to the `HashSet` takes time proportional to the length of the subarray. The length can be up to `O(n)`. Hashing and comparing lists of length `L` takes `O(L)` time. Therefore, the total time complexity is `O(n^3)`.
Space
O(n^3). In the worst-case scenario, we might have `O(n^2)` distinct valid subarrays. If the average length of these subarrays is `O(n)`, the total space required to store them in the `HashSet` would be `O(n^2 * n) = O(n^3)`.
Trade-offs
Pros
Simple to understand and implement.
Correctly solves the problem by exhaustively checking all possibilities.
Cons
The time complexity of
O(n^3)can be too slow if the constraints onnare large.The space complexity of
O(n^3)is high, potentially leading to memory issues.
This approach improves upon the brute-force method by using a more efficient technique to identify distinct subarrays. Instead of storing entire subarray lists in a set, we store their hash values. A rolling hash algorithm allows us to compute the hash of a new subarray (by extending an existing one by one element) in O(1) time. This reduces the overall time complexity from O(n^3) to O(n^2).
Algorithm
- Initialize a
HashSetto store unique hash values of valid subarrays. - Select a suitable base
B(a prime larger than max element value) and a large prime modulusMfor the rolling hash function. - Iterate through all possible start indices
ifrom0ton-1. - For each
i, reset thedivisibleCountto 0 and thecurrentHashto 0. - Iterate through all end indices
jfromiton-1. - In
O(1)time, update the hash for the subarraynums[i...j]using the hash ofnums[i...j-1]with the formula:newHash = (oldHash * B + nums[j]) % M. - Update the count of elements divisible by
p. - If the count is valid (<=
k), add the computed hash to the set. - If the count exceeds
k, break the inner loop. - The final answer is the size of the hash set.
Walkthrough
A rolling hash function computes a hash value for a sequence of items. For a subarray s = [n_1, n_2, ..., n_L], a polynomial rolling hash can be defined as H(s) = (n_1*B^(L-1) + n_2*B^(L-2) + ... + n_L) mod M, where B is a base and M is a large prime modulus to prevent overflow.
The key insight is that when we extend a subarray nums[i...j-1] to nums[i...j], the new hash can be calculated from the old hash in constant time: hash_new = (hash_old * B + nums[j]) mod M. We iterate through all subarrays, calculate their hashes, and store them in a set if they are valid. The size of the set gives the number of distinct valid subarrays.
import java.util.HashSet;import java.util.Set; class Solution { public int countDistinct(int[] nums, int k, int p) { int n = nums.length; Set<Long> distinctHashes = new HashSet<>(); // A prime base larger than the max value of nums[i] (200) long base = 201; // A large prime modulus to reduce collisions long mod = 1_000_000_007; for (int i = 0; i < n; i++) { long currentHash = 0; int divisibleCount = 0; for (int j = i; j < n; j++) { // Update rolling hash currentHash = (currentHash * base + nums[j]) % mod; if (nums[j] % p == 0) { divisibleCount++; } if (divisibleCount <= k) { distinctHashes.add(currentHash); } else { break; } } } return distinctHashes.size(); }}Complexity
Time
O(n^2). The two nested loops give `O(n^2)` iterations. Inside the inner loop, all operations (hash calculation, set insertion, condition check) take `O(1)` average time.
Space
O(n^2). In the worst case, all `O(n^2)` subarrays are valid and distinct, so the `HashSet` will store `O(n^2)` hash values. Each hash value (`long`) takes constant space.
Trade-offs
Pros
Significantly faster than the
O(n^3)brute-force approach, withO(n^2)time complexity.Efficient space usage compared to storing full subarrays.
Relatively simple to implement.
Cons
There is a theoretical, albeit very small, possibility of hash collisions, where two different subarrays produce the same hash value. This can be mitigated by using a second hash function with a different base and modulus, but it adds complexity.
This is a highly efficient and robust approach that achieves the optimal time complexity. A Trie (or prefix tree) is used to store all valid subarrays. Each path from the root of the Trie represents a unique subarray. By inserting all valid subarrays into the Trie, we can count the number of unique ones by simply counting how many new nodes we need to create. This method avoids both the O(L) comparison cost of the brute-force approach and the collision risk of hashing.
Algorithm
- Initialize a
Triedata structure with a single root node and acountof distinct subarrays to 0. - Iterate through the input array
numswith indexifrom0ton-1. Thisimarks the start of a subarray. - For each
i, reset adivisibleCountto 0 and set acurrentNodepointer to the Trie's root. - Start a second loop with index
jfromiton-1. Thisjmarks the end of the subarraynums[i...j]. - Check if
nums[j]is divisible bypand updatedivisibleCount. - If
divisibleCountexceedsk, break the inner loop. - Check if
currentNodehas a child for the valuenums[j]. - If not, it signifies a new unique subarray. Create a new
TrieNode, add it as a child, and incrementcount. - Move
currentNodeto the child node corresponding tonums[j]. - After the loops complete, return
count.
Walkthrough
We define a TrieNode class, where each node contains a map to its children (e.g., HashMap<Integer, TrieNode>). This map links a number to the next node in a sequence.
We iterate through the input array nums with an index i to consider all possible starting points for subarrays. For each start i, we begin a traversal from the Trie's root. A second loop with index j extends the subarray. For each element nums[j], we attempt to move down the Trie. If a path for nums[j] doesn't exist from the current node, it means we've found a new, distinct subarray. We then create a new node, increment our distinct subarray counter, and continue the traversal. We also maintain a count of elements divisible by p and stop extending a subarray if this count exceeds k.
import java.util.HashMap;import java.util.Map; class TrieNode { Map<Integer, TrieNode> children = new HashMap<>();} class Solution { public int countDistinct(int[] nums, int k, int p) { int n = nums.length; TrieNode root = new TrieNode(); int count = 0; for (int i = 0; i < n; i++) { TrieNode currentNode = root; int divisibleCount = 0; for (int j = i; j < n; j++) { if (nums[j] % p == 0) { divisibleCount++; } if (divisibleCount > k) { break; } if (!currentNode.children.containsKey(nums[j])) { currentNode.children.put(nums[j], new TrieNode()); // Each new node represents the end of a new distinct subarray count++; } currentNode = currentNode.children.get(nums[j]); } } return count; }}Complexity
Time
O(n^2). We have two nested loops iterating through the start and end points of the subarrays. The operations inside the inner loop (map lookup, potential insertion, and moving to a child node) take, on average, `O(1)` time.
Space
O(n^2). The number of nodes in the Trie is equal to the number of distinct valid subarrays plus the root. In the worst case, there can be `O(n^2)` such subarrays. The total number of nodes is therefore bounded by `O(n^2)`.
Trade-offs
Pros
Optimal time complexity of
O(n^2).Guaranteed correctness, as it does not suffer from hash collisions.
Cleanly models the problem of finding unique prefixes (which correspond to subarrays in this context).
Cons
Can have slightly more implementation overhead and a larger constant factor in runtime compared to the rolling hash approach due to object creation and map operations.
Solutions
Solution
class Solution {public int countDistinct(int[] nums, int k, int p) { int n = nums.length; Set<String> s = new HashSet<>(); for (int i = 0; i < n; ++i) { int cnt = 0; String t = ""; for (int j = i; j < n; ++j) { if (nums[j] % p == 0 && ++cnt > k) { break; } t += nums[j] + ","; s.add(t); } } return s.size(); }}class Solution {public: int countDistinct(vector<int> &nums, int k, int p) { unordered_set<string> s; int n = nums.size(); for (int i = 0; i < n; ++i) { int cnt = 0; string t; for (int j = i; j < n; ++j) { if (nums[j] % p == 0 && ++cnt > k) { break; } t += to_string(nums[j]) + ","; s.insert(t); } } return s.size(); }};class Solution: def countDistinct(self, nums: List[int], k: int, p: int) -> int: n = len(nums) s = set() for i in range(n): cnt = 0 for j in range(i, n): cnt += nums[j] % p == 0 if cnt > k: break s . add(tuple(nums[i: j + 1])) return len(s)Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.