# K Divisible Elements Subarrays
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/k-divisible-elements-subarrays)
Canonical: https://scaleengineer.com/dsa/problems/k-divisible-elements-subarrays
**Patterns:** [Rolling Hash](https://scaleengineer.com/dsa/patterns/rolling-hash), [Hash Function](https://scaleengineer.com/dsa/patterns/hash-function), [Enumeration](https://scaleengineer.com/dsa/patterns/enumeration)
**Data structures:** Array, Hash Table, Trie
---
## Problem
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 `i` where `nums1[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 <= 200`
* `1 <= nums[i], p <= 200`
* `1 <= k <= nums.length`

**Follow up:**

Can you solve this problem in O(n2) time complexity?

# Approaches
## Brute-Force Generation with a Set
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.
**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)`.
**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 on `n` are large.; The space complexity of `O(n^3)` is high, potentially leading to memory issues.
### Explanation
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.

```java
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();
    }
}
```
### Algorithm
*   Initialize an empty `HashSet` to store distinct valid subarrays, for example, as `Set<List<Integer>>`.
*   Iterate through all possible start indices `i` of a subarray from `0` to `n-1`.
*   For each `i`, iterate through all possible end indices `j` from `i` to `n-1`.
*   For each subarray `nums[i...j]`, create a temporary list and count the number of elements divisible by `p`.
*   A more optimized way is to build the list and count incrementally. For a fixed `i`, as `j` increases, we append `nums[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 the `HashSet`. The set automatically handles uniqueness.
*   If the count exceeds `k`, we can break the inner loop, as any further extension of the subarray from start `i` will also be invalid.
*   The final answer is the size of the `HashSet`.

## Optimized Checking with Rolling Hash
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)`.
**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.
**Pros:** Significantly faster than the `O(n^3)` brute-force approach, with `O(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.
### Explanation
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.

```java
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();
    }
}
```
### Algorithm
*   Initialize a `HashSet` to store unique hash values of valid subarrays.
*   Select a suitable base `B` (a prime larger than max element value) and a large prime modulus `M` for the rolling hash function.
*   Iterate through all possible start indices `i` from `0` to `n-1`.
*   For each `i`, reset the `divisibleCount` to 0 and the `currentHash` to 0.
*   Iterate through all end indices `j` from `i` to `n-1`.
*   In `O(1)` time, update the hash for the subarray `nums[i...j]` using the hash of `nums[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.

## Trie-based Subarray Storage
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.
**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)`.
**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.
### Explanation
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`.

```java
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;
    }
}
```
### Algorithm
*   Initialize a `Trie` data structure with a single root node and a `count` of distinct subarrays to 0.
*   Iterate through the input array `nums` with index `i` from `0` to `n-1`. This `i` marks the start of a subarray.
*   For each `i`, reset a `divisibleCount` to 0 and set a `currentNode` pointer to the Trie's root.
*   Start a second loop with index `j` from `i` to `n-1`. This `j` marks the end of the subarray `nums[i...j]`.
*   Check if `nums[j]` is divisible by `p` and update `divisibleCount`.
*   If `divisibleCount` exceeds `k`, break the inner loop.
*   Check if `currentNode` has a child for the value `nums[j]`.
*   If not, it signifies a new unique subarray. Create a new `TrieNode`, add it as a child, and increment `count`.
*   Move `currentNode` to the child node corresponding to `nums[j]`.
*   After the loops complete, return `count`.

# Solutions
### Java

```java
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();
  }
}

```

### CPP

```cpp
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();
  }
};

```

### Python

```python
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)

```