HSAT 기출 3회: 플레이페어

import sys

message = sys.stdin.readline().rstrip('\n')
key = sys.stdin.readline().rstrip('\n')

alphabet = "ABCDEFGHIKLMNOPQRSTUVWXYZ"

# 좌표값과 테이블 입력
coord = {i: 0 for i in alphabet}
table = [[' '] * 5 for _ in range(5)]

idx = 0
for letter in key + alphabet:
    if coord[letter]:
        continue
    x, y = idx // 5, idx % 5
    table[x][y] = letter
    coord[letter] = (x, y)
    idx += 1

    if idx == 25:
        break

# 암호화
idx = 0
cipher = ""

while idx < len(message):
    m1 = message[idx]
    try:
        m2 = message[idx + 1]
        idx += 2
    except IndexError:
        m2 = 'X'
        idx += 1
    else:
        if m1 == m2:
            m2 = 'Q' if m1 == 'X' else 'X'
            idx -= 1

    x1, y1 = coord[m1]
    x2, y2 = coord[m2]

    if x1 == x2:
        cipher += table[x1][(y1 + 1) % 5]
        cipher += table[x2][(y2 + 1) % 5]
    elif y1 == y2:
        cipher += table[(x1 + 1) % 5][y1]
        cipher += table[(x2 + 1) % 5][y2]
    else:
        cipher += table[x1][y2]
        cipher += table[x2][y1]

print(cipher)