This solution utilizes dynamic programming and bitmasking to generate all possible binary tree structures using the gems as leaves and then computes the numeric value generated by these trees. Special care is taken with mod arithmetic to handle the large numbers produced during the calculation. The core idea is to recursively calculate potential results for each subset of gems represented as bit masks, and combine these results according to the binary tree structure to check if any of these results modulo p equals the target. The bitmask represents which gems are included in a subtree, and recursive calls help construct the left and right subtrees, with memorization (caching) used to avoid redundant calculations.
时间复杂度: O(2^n * n)
空间复杂度: O(2^n)
from functools import cache
from collections import Counter
from math import perm
from typing import List
class Solution:
def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:
# Special handling for edge cases involving prime factors of 10
if p in (2, 5):
if target == 9 % p:
return perm(len(gem) * 2 - 2, len(gem) - 1)
else:
return 0
n = len(gem)
LIMIT = n - n // 3
r10 = pow(10, p - 2, p) # Calculate 10^(p-2) mod p to use in reversal
# Function to iterate through all possible subsets of mask
def iter_mask(mask):
current = 0
while True:
current = ((current | ~mask) + 1) & mask
if current == mask:
break
else:
yield current, (~current & mask)
# Calculate the total 'length' of the number formed by nodes in the subtree defined by mask
@cache
def get_length(mask):
size = -2
for i, g in enumerate(gem):
if (1 << i) & mask:
size += 4 + len(str(g))
return size
# Recursively find all possible numbers formed by the tree rooted at a mask
@cache
def find_all(mask):
if mask.bit_count() == 1:
return [int("1" + str(gem[mask.bit_length() - 1]) + "9") % p]
result = []
for left, right in iter_mask(mask):
left_length = get_length(left)
right_length = get_length(right)
base = pow(10, left_length + right_length + 1, p)
mul = pow(10, right_length + 1, p)
for lr in find_all(left):
lm = lr * mul + base + 9
for rr in find_all(right):
result.append((lm + rr * 10) % p)
return result
# Use Counter to count occurrences of each modulo result
@cache
def find_all_counter(mask):
return Counter(find_all(mask))
# Main function to find how many setups achieve the target modulo
def find_target(mask, target_list):
if mask.bit_count() <= LIMIT:
return sum(find_all_counter(mask).get(t, 0) for t in target_list)
count = 0
for left, right in iter_mask(mask):
left_length = get_length(left)
right_length = get_length(right)
base = pow(10, left_length + right_length + 1, p)
if left.bit_count() <= right.bit_count():
mul = pow(10, right_length + 1, p)
rr_list = []
for lr in find_all(left):
lm = lr * mul + base + 9
rr_list.extend(((t - lm) * r10) % p for t in target_list)
count += find_target(right, rr_list)
else:
mulr = pow(r10, right_length + 1, p)
lr_list = []
for rr in find_all(right):
rm = base + 9 + rr * 10
lr_list.extend(((t - rm) * mulr) % p for t in target_list)
count += find_target(left, lr_list)
return count
return find_target((1 << n) - 1, [target])
sol = Solution()
print(sol.treeOfInfiniteSouls([1, 2, 3], 100, 12319))