2020 KAKAO BLIND RECRUITMENT - 자물쇠와 열쇠 (Level 3)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import numpy as np
from itertools import product


def check(key, lock):
    M, N = len(key), len(lock)
    for s_r, s_c in product(range(N-M+1), range(N-M+1)):
        new_lock = lock.copy()
        new_lock[s_r:s_r+M, s_c:s_c+M] += key
        if np.all(new_lock == 1):
            return True
    return False

def move(key, vals):
    lr, ud = vals

    key = np.roll(key, -ud, axis=0)
    if ud > 0:
        key[-ud:, :] = 0
    elif ud < 0:
        key[:-ud, :] = 0


    key = np.roll(key, lr, axis=1)
    if lr > 0:
        key[:, :lr] = 0
    elif lr < 0:
        key[:, lr:] = 0
    return key

def solution(key, lock):
    key, lock = np.array(key), np.array(lock)
    M = len(key)

    for rot in range(4):
        key_rot = np.rot90(key, rot)
        for vals in product(range(-M+1, M), range(-M+1, M)):
            key_rot_move = move(key_rot, vals)
            if check(key_rot_move, lock):
                return True
    return False