[프로그래머스/Kotlin]Level3 - 양과 늑대

프로그래머스

문제

문제 설명

2진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있습니다. 이 초원의 루트 노드에서 출발하여 각 노드를 돌아다니며 양을 모으려 합니다. 각 노드를 방문할 때 마다 해당 노드에 있던 양과 늑대가 당신을 따라오게 됩니다. 이때, 늑대는 양을 잡아먹을 기회를 노리고 있으며, 당신이 모은 양의 수보다 늑대의 수가 같거나 더 많아지면 바로 모든 양을 잡아먹어 버립니다. 당신은 중간에 양이 늑대에게 잡아먹히지 않도록 하면서 최대한 많은 수의 양을 모아서 다시 루트 노드로 돌아오려 합니다.

예를 들어, 위 그림의 경우(루트 노드에는 항상 양이 있습니다) 0번 노드(루트 노드)에서 출발하면 양을 한마리 모을 수 있습니다. 다음으로 1번 노드로 이동하면 당신이 모은 양은 두 마리가 됩니다. 이때, 바로 4번 노드로 이동하면 늑대 한 마리가 당신을 따라오게 됩니다. 아직은 양 2마리, 늑대 1마리로 양이 잡아먹히지 않지만, 이후에 갈 수 있는 아직 방문하지 않은 모든 노드(2, 3, 6, 8번)에는 늑대가 있습니다. 이어서 늑대가 있는 노드로 이동한다면(예를 들어 바로 6번 노드로 이동한다면) 양 2마리, 늑대 2마리가 되어 양이 모두 잡아먹힙니다. 여기서는 0번, 1번 노드를 방문하여 양을 2마리 모은 후, 8번 노드로 이동한 후(양 2마리 늑대 1마리) 이어서 7번, 9번 노드를 방문하면 양 4마리 늑대 1마리가 됩니다. 이제 4번, 6번 노드로 이동하면 양 4마리, 늑대 3마리가 되며, 이제 5번 노드로 이동할 수 있게 됩니다. 따라서 양을 최대 5마리 모을 수 있습니다.

각 노드에 있는 양 또는 늑대에 대한 정보가 담긴 배열 info, 2진 트리의 각 노드들의 연결 관계를 담은 2차원 배열 edges가 매개변수로 주어질 때, 문제에 제시된 조건에 따라 각 노드를 방문하면서 모을 수 있는 양은 최대 몇 마리인지 return 하도록 solution 함수를 완성해주세요.

제한사항

• 2 ≤ info의 길이 ≤ 17
  • info의 원소는 0 또는 1 입니다.
  • info[i]는 i번 노드에 있는 양 또는 늑대를 나타냅니다.
  • 0은 양, 1은 늑대를 의미합니다.
  • info[0]의 값은 항상 0입니다. 즉, 0번 노드(루트 노드)에는 항상 양이 있습니다.
• edges의 세로(행) 길이 = info의 길이 - 1
  • edges의 가로(열) 길이 = 2
  • edges의 각 행은 [부모 노드 번호, 자식 노드 번호] 형태로, 서로 연결된 두 노드를 나타냅니다.
  • 동일한 간선에 대한 정보가 중복해서 주어지지 않습니다.
  • 항상 하나의 이진 트리 형태로 입력이 주어지며, 잘못된 데이터가 주어지는 경우는 없습니다.
  • 0번 노드는 항상 루트 노드입니다.

입출력 예

info edges result

[0,0,1,1,1,0,1,0,1,0,1,1]	[[0,1],[1,2],[1,4],[0,8],[8,7],[9,10],[9,11],[4,3],[6,5],[4,6],[8,9]]	5
[0,1,0,1,1,0,1,0,0,1,0]	[[0,1],[0,2],[1,3],[1,4],[2,5],[2,6],[3,7],[4,8],[6,9],[9,10]]	5

 

입출력 예 설명

입출력 예 #1

문제의 예시와 같습니다.

입출력 예 #2

주어진 입력은 다음 그림과 같습니다.

0번 - 2번 - 5번 - 1번 - 4번 - 8번 - 3번 - 7번 노드 순으로 이동하면 양 5마리 늑대 3마리가 됩니다. 여기서 6번, 9번 노드로 이동하면 양 5마리, 늑대 5마리가 되어 양이 모두 잡아먹히게 됩니다. 따라서 늑대에게 잡아먹히지 않도록 하면서 최대로 모을 수 있는 양은 5마리입니다.

---

제한시간 안내

• 정확성 테스트 : 10초

풀이

나의 풀이법

풀이 접근 과정

그냥 전체 탐색을 해버리는게 속이 편하겠다 싶었다.

그래서 DFS를 이용하여 모든 노드를 탐색하는 경우의 수를 구하도록 만들었다.

그랬더니 4,5번에서 실패가 생겼다.

알고보니 내가 조건을 잘못 넣어줘서 노드가 두 개 밖에 없고, 모두 양인 경우에 값이 오류가 나왔다.

이 경우에만 0 index 값의 양의 수를 세지 못했다.

그래서 해당 부분을 수정하여 작성한 최종 코드는 아래와 같다.

최종 소스코드

class Solution {
    var answer: Int = 0
    lateinit var parent: IntArray
    lateinit var info: IntArray

    fun solution(info: IntArray, edges: Array<IntArray>): Int {
        parent = IntArray(info.size) { 0 }
        this.info = info

        val visited = BooleanArray(info.size) { false }

        edges.forEach { parent[it[1]] = it[0] }
        dfs(0, visited, 0, 0)

        return answer
    }

    fun dfs(index: Int, visited: BooleanArray, sheepCnt: Int, wolfCnt: Int) {
        if (sheepCnt > answer) answer = sheepCnt
        if (visited[index]) return

        visited[index] = true

        for (i in visited.indices) {
            if (!visited[parent[index]]) continue

            if (info[index] == 0) dfs(i, visited, sheepCnt+1, wolfCnt)
            else {
                if (wolfCnt+1 >= sheepCnt) continue
                else dfs(i, visited, sheepCnt, wolfCnt+1)
            }
        }

        visited[index] = false
    }
}

1. 노드들의 부모 노드를 체크할 parent 배열을 선언한다.
2. 노드들의 방문 여부를 체크할 visited 배열을 선언한다.
3. 그리고 dfs로 모든 노드를 완전 탐색한다.
   1. sheepCnt가 answer보다 크면 answer의 값을 엎어친다.
   2. 그리고 현재 index를 방문한 상태이면 return 한다. (0 index의 계산을 위해 a 뒤에 작성하였다.)
   3. 그리고 i를 0부터 탐색하되, i의 부모 노드를 방문하지 않은 상태면 continue 시킨다.
   4. 만약 현재 값이 양이면, 양의 값을 더해 dfs로 보낸다.
   5. 만약 현재 값이 늑대이면 늑대+1이 양보다 크거나 같아지면 continue 시킨다.
   6. 아니면 늑대의 값을 더해 dfs로 보낸다.

Comment

생각보다 어렵지 않게 풀었다.

완전탐색은 당연히 시간초과가 날 줄 알고 시도해 본것이었지만, info의 길이가 17 밖에 되지 않아 가능했던 것 같다.

풀기는 금방 풀었는데 테스트케이스 4,5번에서 실패가 나서 완전 멘붕이었다.

실패한 시간이 0.05mb 밖에 되지 않아 info의 길이가 짧을 것을 예상하고 때려맞춘 테스트케이스가 적중했다.

TestCase: info ( [0, 0] ) , edges ( [[0, 1]] )

for문 안에서 visited[index]를 리턴시켜버리기 때문에, 0번째 인덱스에서만 오류가 났다.

근데 4,5번만 오류가 나고 나머진 잘 돌아갔다 ???

class Solution {
    var answer: Int = 0
    lateinit var parent: IntArray
    lateinit var info: IntArray

    fun solution(info: IntArray, edges: Array<IntArray>): Int {
        parent = IntArray(info.size) { 0 }
        this.info = info

        val visited = BooleanArray(info.size) { false }

        edges.forEach { parent[it[1]] = it[0] }
        dfs(0, visited, 0, 0)

        return answer
    }

    fun dfs(index: Int, visited: BooleanArray, sheepCnt: Int, wolfCnt: Int) {
        if (sheepCnt > answer) answer = sheepCnt

        visited[index] = true

        for (i in 1 until visited.size) {
            if (!visited[parent[index]] || visited[i]) continue // 여기 때문에 !

            if (info[index] == 0) dfs(i, visited, sheepCnt+1, wolfCnt)
            else {
                if (wolfCnt+1 >= sheepCnt) continue
                else dfs(i, visited, sheepCnt, wolfCnt+1)
            }
        }

        visited[index] = false
    }
}