[Python] 1074 - Z

문제

한수는 크기가 2N × 2N인 2차원 배열을 Z모양으로 탐색하려고 한다. 예를 들어, 2×2배열을 왼쪽 위칸, 오른쪽 위칸, 왼쪽 아래칸, 오른쪽 아래칸 순서대로 방문하면 Z모양이다.

N > 1인 경우, 배열을 크기가 2N-1 × 2N-1로 4등분 한 후에 재귀적으로 순서대로 방문한다.

다음 예는 22 × 22 크기의 배열을 방문한 순서이다.

N이 주어졌을 때, r행 c열을 몇 번째로 방문하는지 출력하는 프로그램을 작성하시오.

다음은 N=3일 때의 예이다.

입력

첫째 줄에 정수 N, r, c가 주어진다.

출력

r행 c열을 몇 번째로 방문했는지 출력한다.

해설

재귀함수 이용해서 전체를 탐색하여 값을 구했었는데 시간초과 + 메모리 초과가 났다. 각 블록마다 짤라서 타겟 블록이 도달하기 이전엔 곱해줘 연산을 줄였다.

소스 코드

import sys
sys.setrecursionlimit(10**7)

n, r, c = map(int, input().split())

count = [0]
def z(n, x, y):
    if n == 0:
        if x == r and y == c:
            print(count[0])
            exit()
        count[0] += 1
    else:
        for a, b in [(0,0),(0,2**(n-1)),(2**(n-1),0),(2**(n-1),2**(n-1))]:
            ax, by = x+a, y+b
            if ax <= r <= ax+2**(n-1)-1 and by <= c <= by+2**(n-1)-1:
                z(n-1, x+a, y+b)
                break
            else:
                count[0] += (2**(n-1))**2
z(n, 0, 0)