Skip to content

2261. K Divisible Elements Subarrays 👍

  • Time: $O(n^2)$
  • Space: $O(n^2)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
struct TrieNode {
  unordered_map<int, shared_ptr<TrieNode>> children;
  int count = 0;
};

class Solution {
 public:
  int countDistinct(vector<int>& nums, int k, int p) {
    int ans = 0;
    for (int i = 0; i < nums.size(); ++i)
      insert(root, nums, i, k, p, ans);
    return ans;
  }

 private:
  shared_ptr<TrieNode> root = make_shared<TrieNode>();

  void insert(shared_ptr<TrieNode> node, const vector<int>& nums, int i, int k,
              int p, int& ans) {
    if (i == nums.size() || k - (nums[i] % p == 0) < 0)
      return;
    if (!node->children.contains(nums[i])) {
      node->children[nums[i]] = make_shared<TrieNode>();
      ++ans;
    }
    insert(node->children[nums[i]], nums, i + 1, k - (nums[i] % p == 0), p,
           ans);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class TrieNode {
  public Map<Integer, TrieNode> children = new HashMap<>();
  public int count = 0;
}

class Solution {
  public int countDistinct(int[] nums, int k, int p) {
    for (int i = 0; i < nums.length; ++i)
      insert(root, nums, i, k, p);
    return ans;
  }

  private int ans = 0;
  private TrieNode root = new TrieNode();

  private void insert(TrieNode node, int[] nums, int i, int k, int p) {
    if (i == nums.length || k - (nums[i] % p == 0 ? 1 : 0) < 0)
      return;
    if (!node.children.containsKey(nums[i])) {
      node.children.put(nums[i], new TrieNode());
      ++ans;
    }
    insert(node.children.get(nums[i]), nums, i + 1, k - (nums[i] % p == 0 ? 1 : 0), p);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class TrieNode:
  def __init__(self):
    self.children: dict[int, TrieNode] = {}
    self.count = 0


class Solution:
  def countDistinct(self, nums: list[int], k: int, p: int) -> int:
    ans = 0
    root = TrieNode()

    def insert(node: TrieNode, i: int, k: int):
      nonlocal ans
      if i == len(nums) or k - (nums[i] % p == 0) < 0:
        return
      if nums[i] not in node.children:
        node.children[nums[i]] = TrieNode()
        ans += 1
      insert(node.children[nums[i]], i + 1, k - (nums[i] % p == 0))

    for i in range(len(nums)):
      insert(root, i, k)

    return ans