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

 

 

프로그래머스 양과 늑대 [ 링크 ]

간단한 문제 요약

초원의 루트 노드에서 시작해, 각 노드를 돌아다니며, 양을 모아야 한다. 노드 방문할 때 마다 해당 노드의 양 또는 늑대는 당신을 따라오고, 이때 당신이 모은 양의 수보다 늑대수가 크거나 같으면 양을 전부 잡아먹는다. 최대한 많은 수의 양을 모아서 루트 노드로 돌아오시오!

문제 풀이

트리를 탐색할 때 dfs 또는 bfs를 떠올렸습니다. 여기서는 다르게 접근해야할게.. 노드에서 좌, 우 서브트리가 있을 경우 좌 서브트리로 탐색할 때, 우 서브트리도 탐색할 수 있어야 합니다.

 

0에서 갈 수 있는 child node는 1, 8입니다. 1을 탐색하고, 어떻게 8로도 탐색할 수 있을까요?!

핵심 포인트는 재귀로 다음 노드를 탐색할 때, 해당 노드에서 탐색할 수 있는 노드 list 정보를 전달해주어야 합니다. 노드의 크기가 17로 크지 않고, 늑대를 탐색할 땐 현재 모아둔 양보다 작은지에 대한 조건을 달아야 합니다.

코드

func makeTree(edges: [[Int]]) -> Int {
  let factory = { (p: inout [Int: [Int]], n: [Int]) in
    p[n[0], default: []].append(n[1))
  }
  return edges.reduce(into: [Int:[Int]](), factory)
}

func solution(_ info: [Int], _ edges: [[Int]]) -> Int {
  let tree = makeTree(edges)
  
  func maxSheep(nodes: [Int], sheep: Int, wolf: Int) {
    if nodes.isEmpty { return sheep }
    var maxSheepCount = sheep
    for node in nodes {
      let isSheep = info[node] == 0
      let nextNodes = nodes.filter{$0 != node }
      if isSheep {
        maxSheepCount = max(maxSheepCount, maxSheep(nodes: nextNodes, sheep: sheep+1, wolf: wolf)
      } else if sheep > wolf+1 {
        maxSheepCount = max(maxSheepCount, maxSheep(nodes: nextNodes, sheep: sheep, wolf: wolf+1)
      }
    }
    return maxSheepCount
  }
  return maxSheep(nodes: tree[0], sheep: 1, wolf: 0)
}