LeetCode 823. Binary Trees With Factors

문제

https://leetcode.com/problems/binary-trees-with-factors/

걸린 시간

- 실패

풀이

C++

typedef long long ll;

class Solution {
public:
    int n;
    int mod = 1e9+7;
    map<int, int> cache;
    int dp(int root, set<int>& s){
        // 메모이제이션
        int& ret = cache[root];
        if(ret != 1) return ret;
        // 재귀 호출
        for(auto x : s){
            if(root%x == 0 && s.find(root/x) != s.end()){
                ret += (ll)dp(x, s)*dp(root/x, s)%mod;
                ret %= mod;
            }
        }
        return ret;
    }
    int numFactoredBinaryTrees(vector<int>& arr) {
        n = arr.size();
        set<int> s(arr.begin(), arr.end());
        int ans = 0;
        for(int x : s) cache[x] = 1;
        for(int x : s){
            ans += dp(x, s);
            ans %= mod;
        }
        return ans;
    }
};

주어진 배열의 고유한 원소를 제한 없이 사용하여, 자식노드 두개의 곱이 부모노드가 되는 이진트리의 총 개수를 찾아야한다.

문제의 조건을 보면 루트 노드만 있는 경우도 수에 포함된다고 한다. (= 가능한 자식 노드가 있지만 표현하지 않은 형태도 경우의 수에 포함된다.)

작은 예제의 모든 경우의 수를 풀어보면 중복되는 계산이 쉽게 보인다. 각 원소를 루트로 하는 경우의 수를 메모이제이션 하면 될 것 같다. 모든 원소는 적어도 자기 자신이 루트인 트리 하나가 존재하므로 초기값은 1 로 설정해 준다.

높이가 2 이상인 트리의 경우 좌, 우 서브트리별로 가능한 경우의 수는 각각 독립 사건이기 때문에 서로 곱해준 값이 해당 루트의 경우의 수가 된다.

dp 를 떠올리는 건 쉬웠지만, root 의 자식노드를 left, right 라고 할 때 곱이 루트노드가 되는 두 자식노드를 찾는 과정에서 left * right = root 인 쌍을 2중 for문으로 찾아 O(n²) 복잡도가 되어 시간초과가 발생했다.

다른 사람의 top-down dp 풀이를 보니 root / left = right 로 형태를 바꾸어 root 가 left 로 나누어 떨어지는지를 확인하고, 나머지 한 쌍(right = root / left)이 원소 배열에 존재하는지를 set 자료구조를 이용해 찾아 복잡도를 O(nlogn) 까지 줄일 수 있었다.