[HackerRank] Forming a Magic Square

0. 들어가며

급하게 코딩테스트가 잡혀서 알고리즘 재활 훈련을 하고 있다.

풀면서 뿌듯했던 문제가 있어 포스팅한다.

1. 문제

Forming a Magic Square | HackerRank

Forming a Magic Square | HackerRank

문제 한 장 요약

문제를 간단하게 요약하자면 아래와 같다.

• 1~9까지의 숫자를 하나씩만 사용하여 가로(3) 세로(3) 대각선(2) 총 8개의 모든 합이 같은 사각형을 magic square라고 부름
• 입력으로 주어진 사각형에서 magic square를 만드는 최소 비용을 구하는 문제

2. 풀이 과정

처음 문제를 읽었을 때는 굉장히 막막했다.

숫자를 바꾸는 순서가 중요한 줄 알고 DFS or 백트래킹을 생각하고 접근했다.

봐도 봐도 해답이 전혀 보이지 않길래 브루트포스 문제구나 싶었다.

 

좀 더 생각해 보니 사각형의 크기가 3*3으로 고정이었다.

1~9를 하나씩 사용하여 사각형 조합을 만드는 모든 경우의 수는 9! = 약 36만 가지.

경우의 수가 굉장히 넉넉해서 모든 magic square를 만들고 비용을 비교하기로 했다.

코드

#!/bin/python3

import math
import os
import random
import re
import sys

#
# Complete the 'formingMagicSquare' function below.
#
# The function is expected to return an INTEGER.
# The function accepts 2D_INTEGER_ARRAY s as parameter.
#

def valid_square(arr):
    s = set()
    for i in range(3):
        s.add(sum(arr[i]))
    for j in range(3):
        res = 0
        for i in range(3):
            res += arr[i][j]
        s.add(res)
    s.add(arr[0][0] + arr[1][1] + arr[2][2])
    s.add(arr[0][2] + arr[1][1] + arr[2][0])
    return len(s) == 1

def square_diff(arr, origin_arr):
    res = 0
    for i in range(3):
        for j in range(3):
            res += abs(arr[i][j] - origin_arr[i][j])
    return res

def make_magic_square(num_set, arr, x, y, origin_arr, res):
    if x == 2 and y == 3:
        if not valid_square(arr):
            return res
        tmp = square_diff(arr, origin_arr)
        return min(res, tmp)
    elif y == 3:
        x += 1
        y = 0

    for n in num_set:
        arr[x][y] = n
        new_set = num_set - {n}
        res = make_magic_square(new_set, arr, x, y + 1, origin_arr, res)
    
    return res

def formingMagicSquare(s):
    num_set = set(range(1, 10))
    arr = [[0] * 3 for _ in range(3)]
    return make_magic_square(num_set, arr, 0, 0, s, 1e9)

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    s = []

    for _ in range(3):
        s.append(list(map(int, input().rstrip().split())))

    result = formingMagicSquare(s)

    fptr.write(str(result) + '\n')

    fptr.close()

 

재귀로 1~9까지의 숫자를 이용해 모든 사각형 조합을 만들었고, 사각형이 magic square인 경우 원본 사각형과 비교하여 비용을 구하였다.

코드가 많이 더럽지만... 풀었으니까 됐다.

3. 느낀 점

오랜만에 구현 문제를 풀어서 풀이 중간중간 실수가 있었지만 무난하게 solve 하였다.

방법만 알면 금방 풀어서 백준으로 치면 실버 1~2 정도? 인 것 같다.

역시 알고리즘이 제일 재밌는 것 같다.