[코드트리/삼성SW역량테스트] 산타의 선물 공장 2 (Python)

Problem : https://www.codetree.ai/training-field/frequent-problems/problems/santa-gift-factory-2

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

Difficulty : Platinum 5

 

Status : Solved

 

Time : 01:03:38

 

---

 

문제 설명 / 입력 및 출력

 

 

자세한 설명은 코드트리 사이트 링크를 참조해 주세요!

 

 

---

 

풀이

 

리스트의 특정 원소 탐색 / 원소 앞뒤 탐색 / 리스트 슬라이싱 및 삽입 등이 빈번하게 일어나는 쿼리를 가지고 있다. 즉 단순히 리스트를 이용하기보다 링크드리스트를 이용하는 것이 시간적으로 이득이다. 링크드리스트를 클래스로 구현하여 생각해보자.

 

 

풀이 코드

import math
import sys
input = sys.stdin.readline

belt_list = list()
Q = int(input())

class Node :
    def __init__(self, idx) :
        self.prev = None
        self.next = None
        self.idx = idx
        
    def present_info(self) :
        a = -1 if self.prev is None else self.prev.idx + 1
        b = -1 if self.next is None else self.next.idx + 1
        print(a + 2*b)

class LinkedList :
    def __init__(self) :
        self.head = None
        self.tail = None
        self.num = 0

    def insert(self, start, end, num) :
        self.num += num
        if self.head == self.tail == None :
            self.head = start
            self.tail = end
            self.head.prev = None
            self.tail.next = None
        else :
            self.head.prev = end
            end.next = self.head
            self.head = start
            self.head.prev = self.tail.next = None
        
    def move_all(self, other_idx) :
        other = belt_list[other_idx]
        if self.head is None :
            print(other.num)
            return
        
        other.insert(self.head, self.tail, self.num)
        print(other.num)
        self.head = self.tail = None
        self.num = 0
        return

    def change_first(self, other_idx) :
        other = belt_list[other_idx]
        
        if other.head is None :
            other_first = None
        else :
            other_first = other.head
            other.head = other_first.next
            if other_first.next is not None :
                other_first.next.prev = None
            else :
                other.tail = None
            other_first.next = None
            other_first.prev = None
            other.num -= 1
        
        if self.head is None :
            self_first = None
        else :
            self_first = self.head
            self.head = self_first.next
            if self_first.next is not None :
                self_first.next.prev = None
            else :
                self.tail = None
            self_first.next = None
            self_first.prev = None
            self.num -= 1
        
        if self_first is not None :
            other.insert(self_first, self_first, 1)
        if other_first is not None :
            self.insert(other_first, other_first, 1)
        
        print(other.num)

    def move_half(self, other_idx) :
        other = belt_list[other_idx]
        if self.num < 2 :
            print(other.num)
            return
        
        half_num = math.floor(self.num / 2)
        start = cur = self.head
        for _ in range(half_num-1) :
            cur = cur.next
        end = cur

        self.head = end.next
        end.next.prev = None
        end.next = None

        self.num -= half_num
        other.insert(start, end, half_num)
        print(other.num)

    def belt_info(self) :
        if self.head is None :
            print(-3)
            return 
        a = self.head.idx + 1
        b = self.tail.idx + 1
        c = self.num
        print(a + 2*b + 3*c)

def init() :
    global belt_list, present_list, N, M

    _, N, M, *lst = map(int, input().split())
    _belt_list = [list() for _ in range(N)]
    present_list = list()
    for i in range(M) :
        belt_num = lst[i]
        present_list.append(Node(i))
        _belt_list[belt_num-1].append(i)

    belt_list = [LinkedList() for _ in range(N)]
    
    for i in range(N) :
        for present_idx in sorted(_belt_list[i], reverse = True) :
            p = present_list[present_idx]
            belt_list[i].insert(p, p, 1)

init()
for _ in range(Q-1) :
    q, *query = map(int, input().split())
    if q == 200 :
        m_src, m_dst = query
        belt_list[m_src-1].move_all(m_dst-1)
    elif q == 300 :
        m_src, m_dst = query
        belt_list[m_src-1].change_first(m_dst-1)
    elif q == 400 :
        m_src, m_dst = query
        belt_list[m_src-1].move_half(m_dst-1)
    elif q == 500 :
        p_num = query[0]
        present_list[p_num-1].present_info()
    else :
        b_num = query[0]
        belt_list[b_num-1].belt_info()

풀이 완료!