[코드트리] 산타의 선물 공장

ewillwin·2024년 4월 13일
0

Problem Solving (CodeTree)

목록 보기
12/16

import sys
input = sys.stdin.readline

class Box:
    def __init__(self, box_id, box_weight):
        self.box_id = box_id
        self.box_weight = box_weight
        self.prev_box = None
        self.post_box = None

    def set_prev(self, prev_box):
        self.prev_box = prev_box

    def set_post(self, post_box):
        self.post_box = post_box

    def cut_prev(self):
        if self.prev_box == None: return
        self.prev_box.post_box = None
        self.prev_box = None
    
    def take_out(self):
        if self.post_box:
            self.post_box.prev_box = self.prev_box
        if self.prev_box:
            self.prev_box.post_box = self.post_box
        self.post_box = None
        self.prev_box = None

class Belt:
    def __init__(self):
        self.first_box = None
        self.last_box = None
        self.box_dict = dict()
        self.broken = False

    def append_box(self, box):
        self.box_dict[box.box_id] = box

        if not self.first_box:
            self.first_box = box
        
        if self.last_box:
            self.last_box.set_post(box)
            box.set_prev(self.last_box)
        
        self.last_box = box

    def pop_box(self):
        box = self.first_box

        self.first_box = box.post_box
        if self.first_box:
            self.first_box.cut_prev()
        else:
            self.last_box = None

        del self.box_dict[box.box_id]

        return box

class Factory:
    def __init__(self, args):
        N, M, *boxs = args
        cnt = N//M
        self.belts = [Belt() for m in range(M)]

        for idx, belt in enumerate(self.belts):
            for i in range(idx*cnt, (idx+1)*cnt):
                belt.append_box(Box(boxs[i], boxs[i+N]))

    def unload_box(self, max_weight):
        result = 0

        for belt in self.belts:
            if belt.first_box:
                box = belt.pop_box()
                if box.box_weight <= max_weight:
                    result += box.box_weight
                else:
                    belt.append_box(box)
        return result

    def remove_box(self, box_id):
        result = -1

        for belt in self.belts:
            if box_id in belt.box_dict:
                box = belt.box_dict[box_id]
                if belt.first_box == box:
                    belt.first_box = box.post_box
                if belt.last_box == box:
                    belt.last_box = box.prev_box
                
                box.take_out()
                result = box.box_id
                del belt.box_dict[box.box_id]
                break
        return result

    def find_box(self, box_id):
        result = -1

        for belt_id, belt in enumerate(self.belts):
            if box_id in belt.box_dict:
                box = belt.box_dict[box_id]

                if belt.first_box != box:
                    belt.first_box.set_prev(belt.last_box)
                    belt.last_box.set_post(belt.first_box)
                    belt.first_box = box
                    belt.last_box = box.prev_box
                    box.cut_prev()
                result = belt_id+1
                break
        return result

    def break_belt(self, belt_id):
        result = -1
        belt_id -= 1
        if self.belts[belt_id].broken == False:
            result = belt_id+1
            broken_belt = self.belts[belt_id]
            broken_belt.broken = True

            if len(broken_belt.box_dict):
                for idx in range(belt_id, belt_id + len(self.belts)):
                    if self.belts[idx % len(self.belts)].broken == False:
                        found_belt = self.belts[idx % len(self.belts)]

                        found_belt.box_dict.update(broken_belt.box_dict)
                        broken_belt.box_dict = {}

                        found_belt.append_box(broken_belt.first_box)
                        found_belt.last_box = broken_belt.last_box
                        broken_belt.first_box = broken_belt.last_box = None
                        break
        return result

if __name__ == "__main__":
    Q = int(input())

    cmd = list(map(int, input().split()))
    factory = Factory(cmd[1:])

    for q in range(1, Q):
        cmd = list(map(int, input().split()))
        if cmd[0] == 200: #물건 하차
            print(factory.unload_box(cmd[1]))
        elif cmd[0] == 300: #물건 제거
            print(factory.remove_box(cmd[1]))
        elif cmd[0] == 400: #물건 확인
            print(factory.find_box(cmd[1]))
        elif cmd[0] == 500: #벨트 고장
            print(factory.break_belt(cmd[1]))
profile
💼 Software Engineer @ LG Electronics | 🎓 SungKyunKwan Univ. CSE

0개의 댓글