K Divisible Elements Subarrays

MEDIUM

Description

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(n²) time complexity?

Approaches

Checkout 3 different approaches to solve K Divisible Elements Subarrays. Click on different approaches to view the approach and algorithm in detail.

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.

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.

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 Analysis

Time Complexity: 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 Complexity: 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 and Cons

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.

Code Solutions

Checking out 3 solutions in different languages for K Divisible Elements Subarrays. Click on different languages to view the code.

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

Video Solution

Watch the video walkthrough for K Divisible Elements Subarrays

Patterns:

Rolling HashHash FunctionEnumeration

Data Structures:

ArrayHash TableTrie