K Divisible Elements Subarrays

Problem

Given an integer array nums and two integers k and p, return the number ofdistinct 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.

Examples

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?

Solution

Method 1 – Hash Set with Sliding Window

Intuition

We want to count the number of distinct subarrays with at most k elements divisible by p. By using a sliding window and a hash set to store unique subarrays, we can efficiently enumerate all possible subarrays and check the divisibility condition.

Approach

1. Initialize an empty set to store unique subarrays (as tuples).
2. For each starting index i in nums:
   • Initialize a counter for divisible elements.
   • For each ending index j from i to the end:
     • If nums[j] % p == 0, increment the counter.
     • If the counter exceeds k, break.
     • Otherwise, add the subarray nums[i:j+1] (as a tuple) to the set.
3. Return the size of the set.

Code

C++

class Solution {
public:
    int countDistinct(vector<int>& nums, int k, int p) {
        int n = nums.size();
        set<vector<int>> st;
        for (int i = 0; i < n; ++i) {
            int cnt = 0;
            vector<int> cur;
            for (int j = i; j < n; ++j) {
                if (nums[j] % p == 0) ++cnt;
                if (cnt > k) break;
                cur.push_back(nums[j]);
                st.insert(cur);
            }
        }
        return st.size();
    }
};

Go

func countDistinct(nums []int, k, p int) int {
    n := len(nums)
    st := map[string]struct{}{}
    for i := 0; i < n; i++ {
        cnt := 0
        cur := []int{}
        for j := i; j < n; j++ {
            if nums[j]%p == 0 {
                cnt++
            }
            if cnt > k {
                break
            }
            cur = append(cur, nums[j])
            key := fmt.Sprint(cur)
            st[key] = struct{}{}
        }
    }
    return len(st)
}

Java

class Solution {
    public int countDistinct(int[] nums, int k, int p) {
        int n = nums.length;
        Set<List<Integer>> st = new HashSet<>();
        for (int i = 0; i < n; ++i) {
            int cnt = 0;
            List<Integer> cur = new ArrayList<>();
            for (int j = i; j < n; ++j) {
                if (nums[j] % p == 0) ++cnt;
                if (cnt > k) break;
                cur.add(nums[j]);
                st.add(new ArrayList<>(cur));
            }
        }
        return st.size();
    }
}

Kotlin

class Solution {
    fun countDistinct(nums: IntArray, k: Int, p: Int): Int {
        val n = nums.size
        val st = mutableSetOf<List<Int>>()
        for (i in 0 until n) {
            var cnt = 0
            val cur = mutableListOf<Int>()
            for (j in i until n) {
                if (nums[j] % p == 0) cnt++
                if (cnt > k) break
                cur.add(nums[j])
                st.add(cur.toList())
            }
        }
        return st.size
    }
}

Python

def count_distinct(nums: list[int], k: int, p: int) -> int:
    n = len(nums)
    st = set()
    for i in range(n):
        cnt = 0
        cur = []
        for j in range(i, n):
            if nums[j] % p == 0:
                cnt += 1
            if cnt > k:
                break
            cur.append(nums[j])
            st.add(tuple(cur))
    return len(st)

Rust

use std::collections::HashSet;
fn count_distinct(nums: Vec<i32>, k: i32, p: i32) -> i32 {
    let n = nums.len();
    let mut st = HashSet::new();
    for i in 0..n {
        let mut cnt = 0;
        let mut cur = vec![];
        for j in i..n {
            if nums[j] % p == 0 {
                cnt += 1;
            }
            if cnt > k {
                break;
            }
            cur.push(nums[j]);
            st.insert(cur.clone());
        }
    }
    st.len() as i32
}

TypeScript

class Solution {
  countDistinct(nums: number[], k: number, p: number): number {
    const n = nums.length;
    const st = new Set<string>();
    for (let i = 0; i < n; ++i) {
      let cnt = 0;
      const cur: number[] = [];
      for (let j = i; j < n; ++j) {
        if (nums[j] % p === 0) cnt++;
        if (cnt > k) break;
        cur.push(nums[j]);
        st.add(cur.join(","));
      }
    }
    return st.size;
  }
}

Complexity

• ⏰ Time complexity: O(n^2 * m) — n is the length of nums, m is the average subarray length (for hashing and storing subarrays).
• 🧺 Space complexity: O(n^2 * m) — For storing all unique subarrays.