[Javascript] 프로그래머스 - 모두 0으로 만들기

문제 설명

각 정점에 가중치가 부여된 트리가 주어질 때, 다음 연산을 이용하여 모든 정점의 가중치를 0으로 만들고자 합니다.

• 임의의 연결된 두 정점을 선택하여 한쪽의 가중치를 1 증가시키고, 다른 한쪽의 가중치를 1 감소시킵니다.

모든 트리가 이 연산을 통해 가중치를 0으로 만들 수 있는 것은 아닙니다. 주어진 트리가 가능 여부를 판별하고, 가능하다면 최소한의 연산 횟수를 구해야 합니다.

---

제한 사항

• a의 길이는 2 이상 300,000 이하입니다.
• a[i]의 값은 -1,000,000 이상 1,000,000 이하입니다.
• edges의 행 개수는 (a의 길이 - 1)이며, 트리의 간선 정보를 나타냅니다.
• edges는 항상 트리의 형태를 유지합니다.

---

해결 방법

알고리즘: 트리 DP, DFS 후위 순회

1. 루트 노드 선정
   • 리프 노드를 루트로 설정하여 트리 탐색을 진행합니다.
2. 트리 구조 생성
   • 인접 리스트를 활용하여 트리를 표현합니다.
3. DFS를 통한 후위 순회 진행
   • DFS를 이용하여 리프 노드부터 순회하며 가중치를 0으로 만드는 연산을 수행합니다.
   • 각 노드에서 부모 노드로 가중치를 전달하며 최소 연산 횟수를 계산합니다.
4. 결과 판별
   • 루트 노드에서 최종 가중치가 0이 아니라면 불가능한 경우이므로 -1을 반환합니다.
   • 가능하다면 최소 연산 횟수를 반환합니다.

---

시간 복잡도

• 트리는 N-1개의 간선을 가지므로, DFS를 이용한 탐색은 O(N)의 시간 복잡도를 가집니다.
• 전체 알고리즘의 시간 복잡도는 O(N)으로 해결할 수 있습니다.

---

구현 코드

function getRootNode(edges) {
  const node = new Array(edges.length + 1).fill(0);

  edges.forEach(([a, b]) => {
    ++node[a];
    ++node[b];
  });

  for (let i = 0; i < node.length; ++i) {
    if (node[i] === 1) {
      return i;
    }
  }
}

function getDist(edges) {
  const dist = Array.from({ length: edges.length + 1 }, () => []);

  edges.forEach(([a, b]) => {
    dist[a].push(b);
    dist[b].push(a);
  });

  return dist;
}

function dfs(start, dist, a) {
  const visit = Array(a.length).fill(false);
  const stack = [[start, null]];
  const order = [];

  while (stack.length) {
    const [node, parent] = stack.pop();

    if (visit[node]) continue;

    visit[node] = true;
    order.push([node, parent]);

    for (const index of dist[node]) {
      if (!visit[index]) stack.push([index, node]);
    }
  }

  const result = Array(a.length).fill(0n);
  const childCount = Array(a.length).fill(0n);

  for (let i = order.length - 1; i >= 0; i--) {
    const [node, parent] = order[i];
    let sum = 0n,
      cnt = 0n;
    for (const index of dist[node]) {
      if (index !== parent) {
        sum += result[index];
        cnt += childCount[index];
      }
    }
    result[node] = BigInt(a[node]) + sum;
    childCount[node] = BigInt(Math.abs(Number(result[node]))) + cnt;
  }

  return [result[start], childCount[start]];
}

function solution(a, edges) {
  const root = getRootNode(edges);
  const dist = getDist(edges);
  const [num, cnt] = dfs(root, dist, a);

  return num === 0n ? cnt : -1;
}