LeetCode 2872. Maximum Number of K-Divisible Components

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

題目大意

在一棵無向樹中，有 n 個節點，0 - n - 1
給定二維整數陣列 edges，其中 edges[i] = [ai, bi] 表示在樹中節點 a_i 和節點 b_i 之間有一條邊

給定整數陣列 values，其中 values[i] 是與第 i 個節點相關聯的值，還有一個整數 k

定義樹之有效分割為：通過刪除樹中的一組邊（可以是空集）使所有的連通分量的值都能被 k 整除，其中連通分量的值是該分量中所有節點的值的總和

回傳任何有效分割中最大的連通分量數量

思考

講到連通分量，那可以考慮看看 DFS

C++ 參考解答：

struct Graph
{
    int numNode;
    vector<vector<int>> adjList;

    Graph(int n, const vector<vector<int>> &edges) : numNode(n), adjList(n)
    {
        for (const auto &edge : edges)
        {
            adjList[edge[0]].push_back(edge[1]);
            adjList[edge[1]].push_back(edge[0]);
        }
    }
};

class Solution
{
public:
    int maxKDivisibleComponents(int n, vector<vector<int>> &edges, vector<int> &values, int k)
    {
        Graph g(n, edges);

        int ans = 0;
        dfs(g, 0, -1, values, k, ans);
        return ans;
    }

private:
    long long dfs(const Graph &g, int u, int prev, const vector<int> &values, const int k, int &ans)
    {
        long long sum = values[u];

        for (const int v : g.adjList[u])
        {
            if (v != prev)
                sum += dfs(g, v, u, values, k, ans);
        }

        if (sum % k == 0)
            ++ans;

        return sum;
    }
};

Go 參考解答：

type Graph struct {
	adjList [][]int
	ans     int
}

func buildGraph(n int, edges [][]int) [][]int {
	graph := make([][]int, n)
	for i := 0; i < n; i++ {
		graph[i] = make([]int, 0, 10)
	}

	for _, edge := range edges {
		u, v := edge[0], edge[1]
		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	return graph
}

func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) int {
	g := Graph{adjList: buildGraph(n, edges), ans: 0}
	g.dfs(0, -1, values, k)
	return g.ans
}

func (g *Graph) dfs(u, p int, values []int, k int) int {
	sum := values[u]

	for _, v := range g.adjList[u] {
		if v != p {
			sum += g.dfs(v, u, values, k)
		}
	}

	if sum%k == 0 {
		g.ans++
	}

	return sum
}