룩(Chess Rook) 문제

팀장님께서 지난 경력직 코딩테스트 문제라며 풀어보라고 주셨다. 알고리즘 공부를 놓은지 6개월 정도가 되었고 알고리즘 문제를 풀 때만 사용하던 파이썬 또한 많이 잊어버려 문법을 검색하며 코딩하였다.

문제는 체스의 룩(Rook) 문제이다.
유명한 문제이고 기존에 풀어봤던 문제지만 풀이과정이 전혀 기억나지 않아 처음 보는 문제를 푸는듯한 느낌으로 풀었다.

이 문제를 풀기 위해 몇가지 방법을 시도하였다.

첫번째, 완전탐색(Exhaustive Search)
이 방법의 경우 N = K일 경우, O(N!)이 되어서 다른 방법으로 넘어갔다.

두번째, 동적 계획법(Dynamic Programming)
이 방법의 경우 첫번째의 완전탐색에 비해 효율적이지만 동일한 연산을 반복한다는 점에서 더욱 효율적으로 개선할 여지가 있다.

def problem(n, k):
  return solution(n, n, k)

def solution(r, c, k):
    if r < k or c < k:
      return 0
    elif k == 0:
      return 1
    else:
      return solution(r-1, c-1, k-1) * c + solution(r-1, c, k)

세번째, 동적 계획법 + 메모이제이션(memoization)
이 방법이 내가 생각한 최종 답이다.
두번째 방법의 동일한 연산을 반복하는 것을 메모이제이션을 이용해 개선하여 효율성을 높였다.

def problem(n, k):
  dp = [[[-1] * (k+1) for _ in range(n+1)] for __ in range(n+1)]
  return solution(n, n, k, dp)

def solution(r, c, k, dp):
  if dp[r][c][k] == -1:
    if r < k or c < k:
      dp[r][c][k] = 0
    elif k == 0:
      dp[r][c][k] = 1
    else:
      dp[r][c][k] = solution(r-1, c-1, k-1, dp) * c + solution(r-1, c, k, dp)

  return dp[r][c][k]

간만에 코딩테스트 문제를 푸니 재미있었고, 많이는 아니더라도 일주일에 세문제 정도 풀며 감을 유지할 수 있도록 해야겠다.