Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 트리의 모든 사과를 수집하는 최소 시간 구하기

정점이 n개로 이루어진 무방향 트리가 있다고 가정해 보겠습니다. 각 정점에는 0부터 n-1까지 번호가 매겨져 있으며, 일부 정점에는 사과가 놓여 있습니다. 트리의 한 간선을 지나갈 때마다 1초가 소요됩니다. 우리가 구해야 할 것은 정점 0에서 출발하여 트리에 있는 모든 사과를 수집한 뒤, 다시 정점 0으로 돌아올 때까지 걸리는 최소 시간(초)입니다.

무방향 트리의 간선 정보는 edges 배열로 주어집니다. edges[i] = [from_i, to_i]는 정점 from_i와 to_i를 연결하는 간선이 존재한다는 의미입니다. 또한 hasApple 배열이 함께 주어지는데, hasApple[i] = true이면 i번 정점에 사과가 있다는 뜻이며, false이면 사과가 없다는 뜻입니다.

예를 들어 n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false, false, true, false, true, true, false]라는 입력이 주어지면 결과는 8이 됩니다.

C++로 트리의 모든 사과를 수집하는 최소 시간 구하기

위 그림에서 빨간색으로 표시된 정점이 사과가 있는 정점입니다. 초록색 화살표는 모든 사과를 수집하는 최적 경로 중 하나를 보여줍니다.

접근 방법

이 문제는 깊이 우선 탐색(DFS)으로 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다. 어떤 서브트리에 사과가 하나라도 존재한다면, 그 서브트리를 탐색하러 들어갔다가 반드시 되돌아와야 하므로 비용은 왕복 기준으로 계산됩니다. 즉, 사과가 있는 서브트리에 진입했다면 내려가는 비용과 올라오는 비용이 모두 발생하기 때문에 간선 비용의 2배를 누적하는 것입니다.

구체적인 해결 단계는 다음과 같습니다.

  • 방문 정보를 저장할 집합(set) visited를 정의합니다.

  • dfs(node, par, a, graph) 함수를 정의합니다. node는 현재 노드, par는 부모 노드, a는 사과 유무 배열, graph는 인접 리스트입니다.

  • temp := 0으로 초기화합니다.

  • graph[node]의 각 원소 x에 대해 다음을 수행합니다.

    • x가 부모 노드 par와 같다면 현재 반복을 건너뜁니다.

    • temp := temp + dfs(x, node, a, graph)

  • ret := ret + temp * 2 (사과가 있는 서브트리를 왕복하는 비용을 누적)

  • a[node] + temp > 0이면 true(1)를, 그렇지 않으면 0을 반환합니다. 이 반환값은 "현재 서브트리에 사과가 존재하는가"를 나타냅니다.

메인 함수에서는 다음과 같이 처리합니다.

  • ret := 0으로 초기화합니다.

  • n개의 리스트를 담는 배열 graph를 정의합니다.

  • i를 0부터 e의 크기 미만일 때까지 1씩 증가시키며 반복합니다.

    • graph[e[i][0]]의 끝에 e[i][1]을 삽입합니다.

    • graph[e[i][1]]의 끝에 e[i][0]을 삽입합니다.

  • dfs(0, -1, a, graph)를 호출합니다.

  • ret을 반환합니다.

예제 코드

더 나은 이해를 위해 다음 구현 코드를 살펴보겠습니다.

#include <bits/stdc++.h>
using namespace std;
const int N = 1e6;
class Solution {
public:
   set<int> visited;
   int ret;
   int dfs(int node, int par, vector<bool>& a, vector<int> graph[]){
      int temp = 0;
      for (int x : graph[node]) {
         if (x == par)
            continue;
         temp += dfs(x, node, a, graph);
      }
      ret += temp * 2;
      return a[node] + temp > 0;
   }
   int minTime(int n, vector<vector<int> >& e, vector<bool>& a){
      ret = 0;
      vector<int> graph[n];
      for (int i = 0; i < e.size(); i++) {
         graph[e[i][0]].push_back(e[i][1]);
         graph[e[i][1]].push_back(e[i][0]);
      }
      dfs(0, -1, a, graph);
      return ret;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{0,1},{0,2},{1,4},{1,5},{2,3},{2,6}};
   vector<bool> v1 = {false,false,true,false,true,true,false};
   cout << (ob.minTime(7,v, v1));
}

입력

7, {{0,1},{0,2},{1,4},{1,5},{2,3},{2,6}},
{false,false,true,false,true,true,false}

출력

8