본문 바로가기
PS/BOJ

[자바] 백준 25516 - 거리가 k이하인 트리 노드에서 사과 수확하기 (java)

by Nahwasa 2022. 9. 2.

 문제 : boj25516


 

필요 알고리즘 개념

  • bfs(너비 우선 탐색), dfs(깊이 우선 탐색)
    • 기본형태의 탐색문제이다. 이 문제의 경우 거리를 따져야 하므로 일반적으로 생각했을 때 bfs가 기본이지만, dfs로도 가능하다. 아무튼 둘 중 하나만 잘 구현할 줄 안다면 풀 수 있다.

※ 제 코드에서 왜 main 함수에 로직을 직접 작성하지 않았는지, 왜 Scanner를 쓰지 않고 BufferedReader를 사용했는지 등에 대해서는 '자바로 백준 풀 때의 팁 및 주의점' 글을 참고해주세요. 백준을 자바로 풀어보려고 시작하시는 분이나, 백준에서 자바로 풀 때의 팁을 원하시는 분들도 보시는걸 추천드립니다.

 


 

풀이

  딱히 뭐 설명이 필요없는 bfs 기본형태의 문제이다. bfs 혹은 dfs에 대해 모른다면 bfs 글을 읽어보자.

 

  그냥 bfs로 하면 심심하니, 내 경우엔 dfs로 풀어봤다. dfs는 재귀로 풀면 되서 짜기 편해서 좋다. 거리가 있는데 dfs로 어떻게 푸냐면, 그냥 코드상의 cnt처럼 현재 몇번째까지 왔는지만 파악해주면 된다. 간선도 양방향 간선이 아니라 부모에서 자식쪽으로 그어주면 되니 방문체크조차 필요없어서 그냥 루트부터 dfs 진행하면서 현재 어디까지 내려왔는지 트리의 깊이를 cnt로 파악해주고, k를 넘었다면 거기서 멈춰주면 된다.

 


 

코드 : github

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Main {
    int n, k;
    ArrayList<Integer>[] edges;
    int[] ex;

    private int dfs(int node, int cnt) {
        if (cnt == k+1) return 0;
        int sum = ex[node];
        for (int next : edges[node]) {
            sum += dfs(next, cnt+1);
        }
        return sum;
    }

    private void solution() throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in), 1<<17);
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());
        ex = new int[n];
        edges = new ArrayList[n];
        for (int i = 0; i < n; i++) edges[i] = new ArrayList<>();
        for (int i = 0; i < n-1; i++) {
            st = new StringTokenizer(br.readLine());
            int p = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());
            edges[p].add(c);
        }
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) ex[i] = st.nextToken().charAt(0)-'0';
        System.out.println(dfs(0, 0));
    }
    public static void main(String[] args) throws Exception {
        new Main().solution();
    }
}

댓글