• Time:O(n^2)
• Space:O(n^2)

## C++

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

## JAVA

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

## Python

``````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, num in enumerate(nums):
insert(root, i, k)

return ans
``````