LeetCode每日一题(2261. K Divisible Elements Subarrays)

Given an integer array nums and two integers k and p, return the number of distinct subarrays which have at most k elements 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

---

这题关键在于去重的方法， 也就是如何 hash 一个整数数组。 搞了半天总是超时， 所以去看答案， 发现了一个叫 Rabin-Karp 的算法， 查了一下， 这个算法貌似更多的用在字符串的查找上， 但是这个题也可以用。具体关于这个算法的的讲解大家去查资料吧， 在这里就不过多赘述了。但是有一点需要注意， 就是我们用的这种 rolling hash 方法是有一定的碰撞概率的， 答案里说的是将不同数组长度的哈希值分开保存，从而减少这一概率， 其他介绍 Rabin-Karp 的文章中提到了用尽可能大的质数来进行取模的方法来减少这一概率。我采用的是前者， 后者没有试， 有兴趣的朋友可以试一下。

---

use std::collections::HashSet;

impl Solution {
    pub fn count_distinct(nums: Vec<i32>, k: i32, p: i32) -> i32 {
        let mut sets: Vec<HashSet<i64>> = vec![HashSet::new(); nums.len()];
        let mut ans = 0;
        'outer: for i in 0..nums.len() {
            let mut count = 0;
            let mut hash = 0i64;
            for j in i..nums.len() {
                count += if nums[j] % p == 0 { 1 } else { 0 };
                if count > k {
                    continue 'outer;
                }
                hash = (200 * hash + nums[j] as i64) % (10i64.pow(9) + 7);
                if !sets[j - i].contains(&hash) {
                    sets[j - i].insert(hash);
                    ans += 1;
                }
            }
        }
        ans
    }
}