2867. Count Valid Paths in a Tree

풀이

• 제한 시간 내에 풀지 못했으며 해답을 참고하였다.

• 에라토스테네스의 체는 이번에야 비로소 이해를 좀 한 듯 하다.

• DFS로 풀어야 하는 것은 맞지만 DP를 적용하는 부분에서 전혀 감이 잡히지 않았었다.

• DP 부분은 현재 노드의 상태 값을 하위 노드들의 상태값을 이용해서 만들어야 한다.
• 상태 값은 prime node가 포함되지 않은 경우 curr[0]와 포함된 경우 curr[1] 를 따로 추적해야 한다.
• 현재 노드가 prime이라면
• prime node가 포함된 경우의 수 curr[1]에 하위 노드 중 prime node가 포함 안된 경우의 수 next[0] 만 추가해야 한다.
• 하위 노드 중 prime node가 포함된 경우의 수는 포함시킬 수 없다. 왜냐하면 그렇게 되면 prime node가 2개 이상 포함된 경로가 되기 때문이다.
• 현재 노드가 prime이 아니라면
• prime node가 포함되지 않은 경우의 수 curr[0]에 하위 노드 중 prime node가 포함 안된 경우의 수 next[0] 를 추가한다.
• prime node가 포함된 경우의 수 curr[1] 에는 하위 노드 중 prime node를 포함한 경우의 수 next[1]를 추가한다.

• 제일 이해가 어려웠던 부분은 경로의 수를 계산하는 부분이었는데 아래와 같이 계산할 수 있다.
• 현재 노드가 prime node라면
• 하위 노드 중 prime node가 아닌 경우의 수 next[0] 를 모두 결과값에 더한다.
• 그리고 하위 노드 중 2개를 선택 (combination) 하여 각각 prime node가 아닌 경우의 수 next[0] 를 구하여 곱한 뒤 결과값에 더한다.
• prime node를 포함하지 않는 두 경로 + 현재 노드 (prime)를 합한 것도 새로운 경로가 된다.
• 현재 노드가 prime node가 아니라면
• 하위 노드 중 prime node인 경우의 수 next[1] 를 모두 결과값에 더한다.
• 그리고 하위 노드 중 2개를 선택 (combination)하여 한쪽은 prime node를 포함한 경우 next[0], 다른쪽은 prime node를 포함하지 않은 경우 next[1]를 선택하여 곱한 뒤 결과값에 더한다.
• prime node를 포함하는 경로 + 현재 노드 (not prime) + prime node를 포함하지 않는 경로를 합한 것도 새로운 경로가 된다.
• 근데 위의 내용을 그대로 구현하게 되면 timeout이 발생하게 된다.
• 위의 내용을 있는 그대로 구현한것은 하단에 추가로 붙여두었다.

class Solution {
public:
    vector<bool> get_prime_table(int n) {
        vector<bool> result(n + 1);
        fill(result.begin(), result.end(), true);
        result[0] = result[1] = false;
        int p = 2;
        while (p <= n) {
            if (!result[p]) ++p;
            int c = p + p;
            while (c <= n) {
                result[c] = false;
                c += p;
            }
            ++p;
        }
        return result;
    }

    array<long long, 2> dfs(
        unordered_map<int, vector<int>>& adj,
        vector<bool>& is_prime,
        int node, int parent, long long& result) {
        array<long long, 2> curr = {};
        auto prime = is_prime[node]? 1 : 0;
        ++curr[prime];
        for (auto child : adj[node]) {
            if (child == parent) continue;
            auto next = dfs(adj, is_prime, child, node, result);
            // non prime cases of others x prime cases of next
            result += curr[0] * next[1];
            // prime cases of others x non prime cases of next
            result += curr[1] * next[0];

            // add non prime cases
            curr[prime] += next[0];
            // count prime cases if current is not prime
            if (!prime) curr[1] += next[1];
        }
        return curr;
    }
    
    long long countPaths(int n, vector<vector<int>>& edges) {
        auto is_prime = get_prime_table(n);
        unordered_map<int, vector<int>> adj;
        for (auto& edge : edges) {
            adj[edge[0]].push_back(edge[1]);
            adj[edge[1]].push_back(edge[0]);
        }
        long long result = 0;
        dfs(adj, is_prime, 1, -1, result);
        return result;
    }
};

DFS를 설명 그대로 구현하면?

• dfs 부분을 설명 그대로 구현한 것은 아래와 같다.
• 하위 노드 2개의 조합에 따라 계산하는 부분이 추가되었고 이로 인해서 시간 제한에 걸리게 된다.
• 최악의 경우 하위 노드가 n에 근접한 숫자가 되면 만큼의 연산을 수행해야 하기 때문이다.

array<long long, 2> dfs(
        unordered_map<int, vector<int>>& adj,
        vector<bool>& is_prime,
        int node, int parent, long long& result) {
        array<long long, 2> curr = {};
        auto prime = is_prime[node]? 1 : 0;
        ++curr[prime];
        // tracking current states
        vector<array<long long, 2>> nexts;
        for (auto child : adj[node]) {
            if (child == parent) continue;
            auto next = dfs(adj, is_prime, child, node, result);
            nexts.push_back(next);
            // add non prime cases
            curr[prime] += next[0];
            // count prime cases if current is not prime
            if (!prime) curr[1] += next[1];
        }
        // calculate result (n^2 time complexity!!!)
        int M = nexts.size();
        if (prime) {
            for (auto& next : nexts) {
                result += next[0];
            }
            for (int i = 0; i < M - 1; ++i) {
                for (int j = i + 1; j < M; ++j) {
                    result += nexts[i][0] * nexts[j][0];
                }
            }
        } else {
            for (auto& next : nexts) {
                result += next[1];
            }
            for (int i = 0; i < M - 1; ++i) {
                for (int j = i + 1; j < M; ++j) {
                    result += nexts[i][0] * nexts[j][1];
                    result += nexts[i][1] * nexts[j][0];
                }
            }
        }
        return curr;
    }

최적화 하기

• 시간 초과를 발생시키는 이유는 두 하위 노드의 값을 곱해서 더하는 것의 시간 복잡도가 이기 때문이다. 따라서 이 부분을 최적화한다.

• 어떤 집합 가 있을 때 집합에서 2개의 숫자를 골라서 곱한 값의 합은 다음과 같이 구할 수 있다.

• 따라서 각 항목의 합을 추적하는 변수 sum 와 현재 값 a[i] 을 곱하여 결과에 더하면 result += sum * a[i] 2개의 숫자를 골라서 곱한 값의 합이 된다.

array<long long, 2> dfs(
        unordered_map<int, vector<int>>& adj,
        vector<bool>& is_prime,
        int node, int parent, long long& result) {
        array<long long, 2> curr = {};
        auto prime = is_prime[node]? 1 : 0;
        ++curr[prime];
        // tracking current states
        vector<array<long long, 2>> nexts;
        for (auto child : adj[node]) {
            if (child == parent) continue;
            auto next = dfs(adj, is_prime, child, node, result);
            nexts.push_back(next);
            // add non prime nodes
            curr[prime] += next[0];
            // count prime nodes if current is not prime
            if (!prime) curr[1] += next[1];
        }
        // calculate result
        if (prime) {
            // add non prime cases
            for (auto& next : nexts) {
                result += next[0];
            }
            // add multiplication of two non prime cases
            long long nonPrimeSum = 0;
            for (auto& next : nexts) {
                result += nonPrimeSum * next[0];
                nonPrimeSum += next[0];
            }
        } else {
            // add prime cases
            for (auto& next : nexts) {
                result += next[1];
            }
            // add multiplication of prime cases & non prime cases
            long long nonPrimeSum = 0;
            long long primeSum = 0;
            for (auto& next : nexts) {
                result += nonPrimeSum * next[1];
                result += primeSum * next[0];
                nonPrimeSum += next[0];
                primeSum += next[1];
            }
        }
        return curr;
    }

• 좀 더 최적화 하자면 sum 로직을 아래처럼 합칠 수 있다.

        // calculate result
        int notPrime = prime == 1? 0 : 1;
        for (auto& next : nexts) {
            result += next[notPrime];
        }
        array<long long, 2> sum = {0, 0};
        for (auto& next : nexts) {
            result += sum[0] * next[1];
            result += sum[1] * next[0];
            sum[prime] += next[0];
            if (!prime) sum[1] += next[1];
        }
        return curr;

• 더 최적화 하면 아래와 같다.

        // calculate result
        array<long long, 2> sum = {0, 0};
        ++sum[prime];
        for (auto& next : nexts) {
            result += sum[0] * next[1];
            result += sum[1] * next[0];
            sum[prime] += next[0];
            if (!prime) sum[1] += next[1];
        }
        return curr;

• 여기서 좀 더 최적화하면 nexts를 제거할 수 있다. 제거한 결과는 본 페이지 처음에 있는 솔루션을 참고하면 된다.