프로그래머스 크레인 인형뽑기 게임 (파이썬)

print()로 정답을 제출했던 백준과 달리, 프로그래머스는 함수형 (return)으로 정답을 반환해야 한다.

간단히 문제 정리를 해 보았다.
▶ 문제 링크

함수 명세

기계 내부가 되는 격자 형태는 N칸 X N칸의 형태이다.

• board
  • N개의 element들로 이루어진 list N개가 argument로 입력된다.
  • 행렬을 떠올리면 쉬울것 같다. 처음 입력된 list부터 가장 상위에 있는 list
  • 숫자 == 인형의 종류 (0은 빈칸)
• moves : 인형을 꺼내올 열 번호 (index 값과 상이하게, 첫 열이 0이 아니라 1이다.)
• return 값 : 사라진 인형의 개수

코드 처리

• board에서 moves에 담긴 열의 맨 위 인형을 꺼내 별도의 바구니에 담는다.
• 별도의 바구니 속 최상단 인형의 종류 == 새로 담는 인형의 종류일 경우 두 인형은 사라진다.
• 사라진 인형의 개수를 반환할 것.

일단 한번 함수를 작성해 보았다.

def solution(board, moves):
    result = 0
    stack = [ ] # 크레인으로 꺼낸 인형을 담음.
    
    for col_idx in moves:
        try:
            col_idx -= 1 # 열 값이니 하나 빼줘야 index 계산 가능
            
            # 뽑을 인형이 위치한 좌표 찾기
            row_idx = [el for el in board if el[col_idx] != 0] # col_idx에 유효한 값이 있는 곳
            row_idx = board.index(row_idx[0]) # 맨 위에 있는 list의 index
            
            # 인형 뽑기 처리
            popped_el = board[row_idx][col_idx] # 젤 위에거, 해당 index의 인형
            board[row_idx][col_idx] = 0 # 빈칸 처리
            
            # 바구니에 담기 전 처리
            if len(stack) > 0: # 빈 바구니가 아닐 경우
                
                if stack[-1] == popped_el: # 맨 위 인형이 같으면
                    stack.pop(-1) # 삭제
                    result += 2 # 터진 인형 : 새로 뽑은 거 + 맨 위에 있던 거
                
                else: # 맨 위 인형이 다른 종류 : stack에 쌓기
                    stack.append(popped_el)
                
            else: # 빈 바구니 : stack에 쌓기
                stack.append(popped_el)
                    
        # moves에 있는 열 번호에 인형이 없는 경우
        except IndexError:
            continue
            
    return result

일단 테스트 코드는 통과했다. 채점버튼을 눌러봤다.

이럴수가. 한번에 정답.
사실 코드를 몇 번 고쳤다.

첫 번째 코드에서의 접근

board를 조작하지 않고, list_to_pop 변수를 만들었다.
col_idx가 0이 아니며, 맨 위에 인형이 위치한 row의 list 변수였다.
그렇지만 이렇게 변수를 따로 빼게 될 경우, board에 접근이 어렵기 때문에
for문이 한번 더 돌때 인형을 뺐음에도 그 전과 같은 board 상태일 것이라 판단했다.

IndexError 예외를 뺀 이유

그리고 프로그래머스 예시 출력을 봤을 때, 1열에 2개의 인형이 있음에도 불구하고 3번을 빼길래 넣었다.

번외 : ChatGPT를 이용해 코드 줄이기

코드가 좀 긴 감이 있어 한번 줄여볼까 싶은 마음이 들었다.
위의 내 코드를 주며 축약을 부탁했고, 주석을 좀 정리해서 코드를 다시 써봤다.

def solution(board, moves):
    result = 0
    stack = [ ] # 크레인으로 꺼낸 인형을 담음.
    
    for col_idx in [col - 1 for col in moves]: # col_idx를 미리 빼 준다.

        # row_idx를 구하는 식은 한줄로 축약이 됐다.
        row_idx = next((row_idx for row_idx, row in enumerate(board) if row[col_idx] != 0), None)
        
        # 코드를 풀어보면
        #   for row_idx, row in enumerate(board):
        #       if row[col_idx] != 0:
        #           return row_idx
        # iterable 반환 -> next(iterable, None) 으로 젤 처음값만 도출
        # -> 가장 위쪽에 있는 인형을 뽑을 행(row_idx)이 됨.
        
        # 뺄 수 있는 인형이 없는 경우
        if row_idx is None: # 예외처리로 빼지 않고 좌표를 도출하기 전 row_idx로 검사.
            continue

        # 빼야할 인형의 좌표는 별도의 축약이 없었다.
        popped_el = board[row_idx][col_idx]
        board[row_idx][col_idx] = 0

        if stack and stack[-1] == popped_el: # stack의 None검사가 합쳐졌다.
            stack.pop()
            result += 2
        else:
            stack.append(popped_el)

    return result

코드가 확실히 짧아졌다. 속도 차이도 있을까 싶어 채점버튼도 다시 한번 눌러봤다.

ChatGPT 코드 ▼

위에서 작성한 내 코드 ▽

작은 양으로 추정되는 데이터가 들어왔을 땐 상관 없지만,
많은 양으로 추정되는 데이터가 들어왔을 땐 대략 0.2 ~ 0.45ms의 속도 차이를 보인다.
아마 무시할 수는 없는 수치라 예상된다.

그리고 주석을 지우면서 나름대로 아쉬웠던 점을 적어봤다.

chatGPT 코드를 참고했을 때 생각되어지는 개선점

1. enumerate() 활용
   • 어차피 board.index()로 row_idx를 빼올거
     enumerate를 돌리면 확실히 빠르고 직관적이였을 것 같다.
2. 내장 함수에 대한 이해 : next()
   • next(iterable[, default]) 의 기본 값 설정 부분을 모르고 있었다.