[프로그래머스] Swift - 양과 늑대

https://school.programmers.co.kr/learn/courses/30/lessons/92343

프로그래머스

 

2023.08.16 기준 Level 3

 

알고리즘 공부를 위해 풀고 기록하는 글입니다!
참고만 해주시고 더 좋은 풀이법이 있다면 알려주세요!

 

아이디어

1. 이진 트리의 각 노드에 양 또는 늑대가 존재
2. 노드를 방문하면 그 노드에 있는 동물을 획득
3. 늑대의 수가 양의 수보다 크거나 같아지면 양들 모두 사망
4. 구할 수 있는 양의 최댓값을 구해야 한다.

언뜻 보면 단순한 트리 순회 문제 같지만 1가지 특이한 점이 있습니다.

보통 일반적인 트리 순회는 한 노드를 방문하고 그 노드의 자식 노드들을 쭉 방문하는 방식인데 이 문제는 구할 수 있는 양의 최댓값을 위해서는 노드를 방문했다가 부모나 자식 노드가 아닌 반대쪽 노드로 탐색이 전환될 수 있다는 점입니다.

 

위의 예시에서도 0 ➡️ 1 ➡️ 8 ➡️ 7 ➡️ 9 를 탐색하고 9의 반대 subtree에 있는 4 ➡️ 6 ➡️ 5 로 탐색을 이어가야 정답을 구할 수 있습니다.

 

따라서, dfs나 bfs로 트리 탐색을 진행할 때 다음에 방문할 노드를 정하는 로직에서 자신의 현재 노드 뿐만 아니라 방문할 수 있는 노드 전체를 대상으로 놓고 방문해야 한다는 것이 주요했습니다.

 

여기서 방문할 수 있는 노드 전체를 이해해야 합니다.

 

만약 위의 예시에서 0번 노드를 방문했다면 1번과 8번을 방문할 수 있는 길이 열린 것입니다.

그리고 1번 노드를 방문한 순가 2, 4번 노드로의 길이 열렸습니다.

그 다음 8번 노드를 방문하면 7, 9번 노드로도 길이 열린 것으로 생각할 수 있습니다.

 

방문할 수 있는 노드를 배열로 나타내면

0번 노드 방문 후 -> [0, 1, 8]  (방문 가능한 전체 노드를 의미합니다.)

1번 노드 방문 후 ->  [0, 1, 8, 2, 4]

8번 노드 방문 후 -> [0, 1, 8, 2, 4, 7, 9]

...

이렇게 노드 하나를 방문할 때마다 마치 도로가 뚫리듯 갈 수 있는 곳이 늘어나게 됩니다.

여기서 방문 가능하다는 의미는 8번 노드를 탐색하다가 곧바로 2번 노드로 옮겨 탐색을 이어 나갈 수 있다는 의미입니다.

 

이제 이 아이디어를 이어나가면 이 문제를 해결할 수 있습니다.

 

1. DFS로 탐색
2. 0번 노드부터 탐색 시작
3. 방문 가능한 노드들을 DFS 함수의 파라미터로 넘김
4. 방문 가능한 노드들을 순회하며 전부 방문
5. 3~4 번 작업을 재귀적으로 수행하며 방문한 노드의 동물에 따라 양과 늑대의 수를 기록합니다.
6. 늑대의 수가 양의 수보다 같거나 많아지면 재귀를 종료합니다.
7. 양의 수의 최댓값을 계속 갱신하여 재귀가 모두 종료되면 이 값을 리턴합니다.

 

이 흐름으로 코드를 구현했습니다.

 

여기서 3번 작업이 중요했습니다.

앞서 방문할 수 있는 노드의 의미를 설명했는데 이 때 실제 코드에서는 방문할 노드는 이 배열에서 제거해야 중복 탐색을 막을 수 있습니다.

예를 들어 8번 노드를 방문할 건데 넘겨줄 방문 가능한 노드 배열에 8번 노드가 또 들어있으면 중복 계산이 발생하기 때문입니다.

따라서 자기 자신을 제외한 노드 번호들만 전달하게 되기 때문에 8번 노드를 처음 방문할 때 앞으로 방문 가능한 노드 배열(= 방문 할 노드)는 [2, 4]이고 이후 8번의 길이 열렸기 때문에 자식 노드를 추가한 [2, 4, 7, 9]가 다음 탐색 대상이 됩니다.

 

 

풀이

fileprivate func solution(_ info:[Int], _ edges:[[Int]]) -> Int {
    var tree = [Int: [Int]]()
    
    for edge in edges {
        tree[edge[0], default: []].append(edge[1])
    }
    
    func isSheep(_ node: Int) -> Bool {
        return info[node] == 0
    }

    var result = 0
    
    func dfs(cur: Int, sheepCnt: Int, wolfCnt: Int, possibleNodes: [Int]) {
        var sheepCnt = sheepCnt
        var wolfCnt = wolfCnt
        var possibleNodes = possibleNodes
        
        if isSheep(cur) {
            sheepCnt += 1
        } else {
            wolfCnt += 1
        }
        
        if wolfCnt >= sheepCnt {
            return
        }
        
        result = max(result, sheepCnt)
        
        let children = tree[cur, default: []]
        possibleNodes.append(contentsOf: children)
        
        for node in possibleNodes {
            dfs(cur: node, sheepCnt: sheepCnt, wolfCnt: wolfCnt, possibleNodes: possibleNodes.filter { $0 != node })
        }
    }
    
    dfs(cur: 0, sheepCnt: 0, wolfCnt: 0, possibleNodes: [])
    
    return result
}

• tree 딕셔너리에 edges 데이터를 저장합니다. key는 부모 노드 value는 자식 노드입니다.
• 양이면 sheepCnt를 1 더하고 늑대면 wolfCnt를 1만큼 더합니다.
• wolfCnt >= sheepCnt는 종료 조건입니다. (base case)
• possibleNodes는 방문 가능한 노드들을 담은 배열입니다.
• 방문한 현재 노드의 자식 노를 이 배열에 추가합니다.
• 다음 방문을 위해 dfs 함수를 호출 할 때 현재 방문한 노드를 possibleNodes에서 제거하고 넘겨줍니다.
  • filter로 제거!

 

 

 

마무리

트리 문제는 익숙한 편인데 이렇게 조금 변형되면 아이디어를 찾기 쉽지 않은 것 같습니다.

기본적은 트리 풀이법에서 다음 방문할 노드를 고려하는 것이 중요했습니다.