Problem URL: Binary trees with factors

We will convert the array to array set, so that we can do the lookup in constant time. Then we will check, if a candidate can divide the root without any remainder and root//candidate is available in the array set, then we add it to our answer. Finally, we sum up all the possible answer for every element of the array and mod it by our given mod.

class Solution:
    def numFactoredBinaryTrees(self, arr: List[int]) -> int:
        mod = 10**9+7
        arrSet = set(arr)

        @cache
        def dp(root):
            ans = 1
            for candidate in arrSet:
                if root % candidate == 0 and (root // candidate) in arrSet:
                    ans += dp(candidate) * dp(root // candidate)
                    ans %= mod
            return ans

        return sum(dp(x) for x in arrSet) % mod

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