Softeer 연습문제(3단계) - 플레이페어 암호

Candidate | Softeer Assessment UI

Candidate | Softeer Assessment UI

 

나의 코드

import sys
from collections import deque
import copy
input = sys.stdin.readline

# 메시지 키 가져옴
msg_vl = input().split()[0]
key_vl = input().split()[0]

# 초기 알파벳
alpa = ['A','B','C','D','E','F','G','H','I','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
#print(msg_vl, key_vl)

# map에 들어갈 키 만들어주자
change_key = []
for i in key_vl:
    if i not in change_key:
        change_key.append(i)
        if i in alpa:
            alpa.remove(i)
change_key = change_key+alpa

#map 만들어 주기
argo_map = []
for i in range(0, 21, 5):
    argo_map.append(change_key[i:i+5])
#print(argo_map)

# 메세지 나누기(암호화 하려는 마세지 두 글자씩 나눔)
argo_massage = []
msg_dict = {'O':[], 'X':[]}
queue = deque(msg_vl)
group_msg=''
while queue:
    now_value = queue.popleft()
    if now_value not in msg_dict:
        msg_dict[now_value]=[]

    if not group_msg:
        group_msg+=now_value
    elif group_msg!=now_value:
        group_msg+=now_value
        argo_massage.append(group_msg)
        group_msg=''
    elif group_msg==now_value:
        if group_msg=='X':
            argo_massage.append(group_msg+'Q')
        else:
            argo_massage.append(group_msg+'X')
        group_msg=now_value
if group_msg:
    if group_msg=='X':
        argo_massage.append('XQ')
    else:
        argo_massage.append(group_msg+'X')
#print(argo_massage)

# 메세지에 사용된 문자의 argo_map 상 인덱스를 저장한 딕셔너리
for i in range(len(argo_map)):
    for j in range(len(argo_map[0])):
        if argo_map[i][j] in msg_dict:
            msg_dict[argo_map[i][j]]=[i,j]
#print(msg_dict)

# 두글자 리스트 암호화
def result_key(msg_list, msg_dict):
    answer = []
    #print(msg_list, msg_dict)
    for i in msg_list:
        a,b = i
        a_x,a_y = copy.deepcopy(msg_dict[a])
        b_x,b_y = copy.deepcopy(msg_dict[b])
        # 행이 같으면
        if a_x==b_x:
            a_y=(a_y+1)%(len(argo_map[0]))
            b_y=(b_y+1)%(len(argo_map[0]))
        # 열이 같으면
        elif a_y==b_y:
            a_x=(a_x+1)%(len(argo_map))
            b_x=(b_x+1)%(len(argo_map))
        # 둘다 아니면
        else:
            a_y,b_y = b_y,a_y
        answer.append(argo_map[a_x][a_y])
        answer.append(argo_map[b_x][b_y])
    return ''.join(answer)
    
print(result_key(argo_massage, msg_dict))

 

하지만.. 런타임 에러 최적해 해보자고~!

 

수정된 코드

import sys
from collections import deque

input = sys.stdin.readline

# 메시지와 키 가져오기
msg_vl = input().strip()
key_vl = input().strip()

# 초기 알파벳
alpa = ['A','B','C','D','E','F','G','H','I','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']

# 키에 따라 5x5 암호 표 생성
change_key = []
key_set = set()
for char in key_vl:
    if char not in key_set:
        change_key.append(char)
        key_set.add(char)

for char in alpa:
    if char not in key_set:
        change_key.append(char)

argo_map = [change_key[i:i+5] for i in range(0, 25, 5)]

# 메세지 나누기
argo_message = []
queue = deque(msg_vl)
while queue:
    first = queue.popleft()
    if queue:
        second = queue.popleft()
        if first == second:
            second = 'X' if first != 'X' else 'Q'
            queue.appendleft(first)
    else:
        second = 'X'
    argo_message.append(first + second)

# 문자 위치 저장
msg_dict = {char: (i, j) for i, row in enumerate(argo_map) for j, char in enumerate(row)}

# 암호화 함수
def result_key(msg_list, msg_dict):
    answer = []
    for pair in msg_list:
        a, b = pair
        a_x, a_y = msg_dict[a]
        b_x, b_y = msg_dict[b]
        if a_x == b_x:
            answer.append(argo_map[a_x][(a_y + 1) % 5])
            answer.append(argo_map[b_x][(b_y + 1) % 5])
        elif a_y == b_y:
            answer.append(argo_map[(a_x + 1) % 5][a_y])
            answer.append(argo_map[(b_x + 1) % 5][b_y])
        else:
            answer.append(argo_map[a_x][b_y])
            answer.append(argo_map[b_x][a_y])
    return ''.join(answer)

# 암호화 결과 출력
print(result_key(argo_message, msg_dict))

 

후하후하... 점점 어려워진다...