백준15787 - 기차가 어둠을 헤치고 은하수를 (비트마스킹) Python

https://www.acmicpc.net/problem/15787

해당 문제는 열차마다 20개의 승객 좌석이 주어졌고, 좌석은 승객이 탄 경우 | 타지 않은 경우로 나눌 수 있기 때문에

어제 포스팅한 비트마스킹 풀이로 접근했다.

2023.11.17 - [개발/Python | Java] - Python 문법 & 비트마스킹

 

풀이 과정은 다음과 같다.

1. 각 열차마다 승객의 상태를 확인하기 위한 배열을 선언한다.

  arr = list(map(lambda x : list(map(int, x.split())), data))
  n, _ = arr.pop(0) # n은 열차의 갯수, _은 명령의 갯수
  
  trains = [0] * (n+1)

 

2. 주어진 명령들을 처리하고 해당 배열에 업데이트해나간다.

 for cmd in arr :
    order, train_num = cmd[0], cmd[1]
    temp = None
    # 명령 종류인 order가 2 이하인 경우, 원소 추가 | 제거
    # 2 초과인 경우, 좌측 이동 | 우측 이동이므로 이를 구분하여 다루기 위해 if문 사용
    if 2 >= cmd[0] :
      seat = cmd[2] - 1 # 추가 | 제거될 좌석
      # order가 1인 경우 해당 위치에 추가 | 2인 경우 해당 위치 제거
      temp = trains[train_num] | (1 << seat) if order == 1 else trains[train_num] & ~(1 << seat)
    else :
      # order가 3인 경우 좌측이동 | 4인 경우 우측 이동
      temp = trains[train_num] << 1 if order == 3 else trains[train_num] >> 1
      # 이때, order가 3이며, 비트 길이가 20을 초과하는 경우, 초과된 길이를 잘라내는 과정이 필요
      if temp & (1 << 20) : temp &= ~(1 << 20)
    trains[train_num] = temp

 

3. 타 승객 위치와 중복되지 않은 열차의 갯수만 구한 뒤 출력한다.

  result = 0
  # 승객 위치를 기록하기 위한 딕셔너리 선언
  storage = {}
  for v in trains[1:] :
    if not storage.get(v, False) :
      storage[v] = True
      result += 1
  
  print(result)

 

전체 소스 코드

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()

def solution (data) :
  arr = list(map(lambda x : list(map(int, x.split())), data))
  n, _ = arr.pop(0)
  
  trains = [0] * (n+1)
  
  for cmd in arr :
    order, train_num = cmd[0], cmd[1]
    temp = None
    if 2 >= cmd[0] :
      seat = cmd[2] - 1
      temp = trains[train_num] | (1 << seat) if order == 1 else trains[train_num] & ~(1 << seat)      
    else :
      temp = trains[train_num] << 1 if order == 3 else trains[train_num] >> 1
      if temp & (1 << 20) : temp &= ~(1 << 20)
    trains[train_num] = temp
    
  result = 0
  storage = {}
  for v in trains[1:] :
    if not storage.get(v, False) :
      storage[v] = True
      result += 1
  
  print(result)
  
solution(input_data)