[LeetCode] 2872. Maximum Number of K-Divisible Components

https://leetcode.com/problems/maximum-number-of-k-divisible-components/description/

leetcode - Maximum Number of K-Divisible Components

문제 유형: 트리, dfs, 부분합

문제 난이도: Hard

 

문제

There is an undirected tree with n nodes labeled from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

You are also given a 0-indexed integer array values of length n, where values[i] is the value associated with the ith node, and an integer k.

A valid split of the tree is obtained by removing any set of edges, possibly empty, from the tree such that the resulting components all have values that are divisible by k, where the value of a connected component is the sum of the values of its nodes.

Return the maximum number of components in any valid split.

 

n개의 노드가 0~n-1번호로 되어 무방향 트리를 구성한다. n-1길이의 간선 정보도 주어진다.

정수 k도 받는다. 각 노드의 값이 담긴 n길이의 values배열도 받는다.

 

유효한 분할은 모든 컴포넌트 내의 값의 합이 k의 배수인 경우를 말한다.

최대한 컴포넌트들로 분하여 개수를 구하시오. 반드시 보장되는 테스트케이스만 나옵니다.

 

풀이

만약 리프노드가 k의 배수가 아니라면, 부모 노드와 결합되어야 한다.

아래 서브트리가 독립될 수 없다면, 적어도 부모 노드와 결합되어야 한다.

 

이 성질을 이용해서 dfs를 이용하면 된다. 자식 노드들로부터 누적합을 구해온다.

누적합이 k의 배수라면, 컴포넌트 분리가 된다. 컴포넌트 개수를 1 추가한다.

부모노드에게 자신의 누적합을 올린다. 어차피 k의 배수인 경우라도 상위 서브트리에서의 계산 반영에 영향을 끼치지 않는다.

 

이러한 dfs를 루트 0, 부모 -1로 호출하여 풀어낸다.

 

코드

C++

class Solution {
public:
    int result = 0;
    int k;
    vector<vector<int>> graph;

    long long dfs(int node, int parent, vector<int>& values) {
        long long currentSum = values[node];
        
        for (int neighbor : graph[node]) {
            if (neighbor != parent) {
                currentSum += dfs(neighbor, node, values);
            }
        }
        
        if (currentSum % k == 0) {
            result++;
            return 0;
        }
        
        return currentSum;
    }

    int maxKDivisibleComponents(int n, vector<vector<int>>& edges, vector<int>& values, int k) {
        this->k = k;
        graph.resize(n);
        
        for (const auto& edge : edges) {
            int u = edge[0], v = edge[1];
            graph[u].push_back(v);
            graph[v].push_back(u);
        }
        
        dfs(0, -1, values);
        
        return result;
    }
};