K Divisible Elements Subarrays

medium trie array enumeration

Problem

Given an array nums and integers k and p, return the number of distinct subarrays that have at most k elements divisible by p. A subarray is a non-empty contiguous slice; two subarrays are distinct if they differ in length or in any position.

Inputnums = [2,3,3,2,2], k = 2, p = 2

Output11

Indices 0, 3, 4 are divisible by 2. There are 11 distinct subarrays with at most 2 such elements; [2,3,3,2,2] is excluded (3 divisible) and repeats like [2] count once.

nums (comma-separated)

def countDistinct(nums, k, p):
    root = {}              # trie root (dict of children)
    count = 0              # distinct subarrays = new trie nodes
    n = len(nums)
    for i in range(n):     # start each subarray at index i
        node = root
        div = 0            # count of elements divisible by p
        for j in range(i, n):
            if nums[j] % p == 0:
                div += 1
            if div > k:    # too many divisible: stop extending
                break
            if nums[j] not in node:
                node[nums[j]] = {}   # new path = new subarray
                count += 1
            node = node[nums[j]]
    return count

function countDistinct(nums, k, p) {
  const root = new Map();   // trie root
  let count = 0;            // distinct subarrays = new nodes
  const n = nums.length;
  for (let i = 0; i < n; i++) {   // start at index i
    let node = root;
    let div = 0;            // elements divisible by p so far
    for (let j = i; j < n; j++) {
      if (nums[j] % p === 0) div++;
      if (div > k) break;   // too many divisible: stop
      if (!node.has(nums[j])) {
        node.set(nums[j], new Map());  // new path = new subarray
        count++;
      }
      node = node.get(nums[j]);
    }
  }
  return count;
}

int countDistinct(int[] nums, int k, int p) {
    Map<Integer, Object> root = new HashMap<>();  // trie root
    int count = 0;          // distinct subarrays = new nodes
    int n = nums.length;
    for (int i = 0; i < n; i++) {   // start at index i
        Map<Integer, Object> node = root;
        int div = 0;        // elements divisible by p so far
        for (int j = i; j < n; j++) {
            if (nums[j] % p == 0) div++;
            if (div > k) break;     // too many divisible: stop
            if (!node.containsKey(nums[j])) {
                node.put(nums[j], new HashMap<Integer, Object>());
                count++;            // new path = new subarray
            }
            node = (Map<Integer, Object>) node.get(nums[j]);
        }
    }
    return count;
}

struct Node { unordered_map<int, Node*> nxt; };

int countDistinct(vector<int>& nums, int k, int p) {
    Node* root = new Node();   // trie root
    int count = 0;             // distinct subarrays = new nodes
    int n = nums.size();
    for (int i = 0; i < n; i++) {   // start at index i
        Node* node = root;
        int div = 0;           // elements divisible by p so far
        for (int j = i; j < n; j++) {
            if (nums[j] % p == 0) div++;
            if (div > k) break;        // too many divisible: stop
            if (!node->nxt.count(nums[j])) {
                node->nxt[nums[j]] = new Node();  // new subarray
                count++;
            }
            node = node->nxt[nums[j]];
        }
    }
    return count;
}

Explanation

We must count distinct subarrays, so naive enumeration would over-count repeats like [2] appearing several times. A trie solves the de-duplication elegantly: each distinct subarray corresponds to exactly one root-to-node path, so a brand-new path means a brand-new subarray.

For every start index i we walk forward, descending the trie one value at a time. We keep div, the running count of elements divisible by p in the current window nums[i..j]. The moment div exceeds k we break, because any longer subarray starting at i would also be invalid.

While descending, if the current value is not yet a child of the current trie node we create that child and increment count — this single new node represents the subarray nums[i..j] that has never been seen before. If the child already exists, this prefix (and hence this subarray) was produced by an earlier start index, so we do not count it again.

In Example 1, starting at index 0 we insert [2], [2,3], [2,3,3], [2,3,3,2] (4 nodes), then div would hit 3 at [2,3,3,2,2] so we stop. Repeating from every start index and only counting freshly created nodes yields exactly 11.

Because each (start, end) pair is visited at most once and a trie step is O(1), the whole scan is O(n²) — matching the follow-up — and the trie stores at most one node per distinct prefix.

Time: O(n²) Space: O(n²)