코드트리 산타의 선물 공장 2 (파이썬)

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

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

문제대로 벨트에 있는 모든 박스 위치를 매번 반복문으로 전부 갱신하면 시간초과가 뜬다

답을 구하는데 필요한 정보는 벨트위의 박스개수, 박스의 앞/뒤 박스 번호, 벨트위의 맨 앞/뒤 박스 번호이므로 해당 변수만 갱신해나가면 해결할 수 있다

 

1. 각 박스 번호를 인덱스로 하는 리스트 prev, nxt를 만든다

    prev는 박스의 앞, nxt는 박스의 뒷 번호를 넣는다

2. 벨트의 번호를 인덱스로 하는 리스트 front, back, len_conv를 만든다

    front는 맨 앞 박스의 번호 , back은 맨 뒤 박스의 번호, len_conv는 박스의 개수를 넣는다

3. 맨 처음 입력받은 박스를 conv에 순서대로 기록한다

   초기상태에 대한 각 벨트와 박스에 대한 정보들을 저장한다

4. 각 명령에 대해서 구현한다

   박스를 이동시킬 벨트를 a / 박스를 받을 벨트를 b라고 할 때, b의 현재 박스 개수가 0개 일 경우를 따로 고려해야한다

 

4-1) 전부 옮기기

옮기면서 바뀌는 정보들을 수정한다

이 때, b에 현재 박스가 없으면 b의 맨 뒤 박스 번호을 a의 맨 뒤 박스 번호로 수정한다

 

4-2) 앞 물건만 교체

a b 모두 박스가 있을 경우, 어느 한쪽에만 박스가 있을 경우로 나눠서 구현

1) 모두 박스가 있는데, 딱 한개씩만 존재하는 경우 맨 뒤 박스 번호도 수정해준다

2) 어느 한 쪽만 있는데 이동하면서 해당 벨트의 박스가 0개가 되면 맨 뒤 박스도 지워준다

 

4-3) 나누기

옮길 박스의 번호 num을 반복문으로 찾아낸다

해당 번호를 기준으로 박스를 옮긴다

이 때, b에 현재 박스가 없으면 b의 맨 뒤 박스 번호을 num으로 수정한다

import sys, math

input = sys.stdin.readline

q = int(input())
inp = list(map(int, input().split()))
n, m = inp[1], inp[2]

prev, nxt = [0]*(m+1), [0]*(m+1)
front, back, len_conv = [0]*(n+1), [0]*(n+1), [0]*(n+1)

conv = [[] for _ in range(n+1)]
for idx, num in enumerate(inp[3:]):
    conv[num].append(idx+1)

for i in range(n+1):
    if not conv[i]:
        continue

    front[i] = conv[i][0]
    back[i] = conv[i][-1]
    len_conv[i] = len(conv[i])

    if len(conv[i]) == 1:
        continue

    for j in range(len(conv[i])-1):
        prev[conv[i][j+1]] = conv[i][j]
        nxt[conv[i][j]] = conv[i][j+1]

for _ in range(q-1):
    order = list(map(int, input().split()))

    if order[0] == 200:
        src, dst = order[1], order[2]
        if len_conv[src] != 0:
            nxt[back[src]] = front[dst]
            prev[front[dst]] = back[src]
            front[dst], front[src] = front[src], 0
            if len_conv[dst] == 0:
                back[dst] = back[src]
            back[src] = 0
            len_conv[dst] += len_conv[src]
            len_conv[src] = 0
        print(len_conv[dst])

    elif order[0] == 300:
        src, dst = order[1], order[2]
        if len_conv[src] != 0 and len_conv[dst] != 0:
            prev[nxt[front[src]]], prev[nxt[front[dst]]] = front[dst], front[src]
            nxt[front[src]], nxt[front[dst]] = nxt[front[dst]], nxt[front[src]]
            front[src], front[dst] = front[dst], front[src]
            if len_conv[src] == 1:
                back[src] = front[src]
            if len_conv[dst] == 1:
                back[dst] = front[dst]

        elif len_conv[src] != 0 and len_conv[dst] == 0:
            prev[nxt[front[src]]] = 0
            front[src], front[dst] = nxt[front[src]], front[src]
            nxt[front[dst]] = 0
            back[dst] = front[dst]
            len_conv[dst] += 1
            len_conv[src] -= 1
            if len_conv[src] == 0:
                back[src] = 0

        elif len_conv[src] == 0 and len_conv[dst] != 0:
            prev[nxt[front[dst]]] = 0
            front[dst], front[src] = nxt[front[dst]], front[dst]
            nxt[front[src]] = 0
            back[src] = front[src]
            len_conv[src] += 1
            len_conv[dst] -= 1
            if len_conv[dst] == 0:
                back[dst] = 0

        print(len_conv[dst])

    elif order[0] == 400:
        src, dst = order[1], order[2]
        if len_conv[src] > 1:
            idx = math.floor(len_conv[src]/2)
            if idx == 1:
                num = front[src]
            else:
                before = front[src]
                for _ in range(idx-1):
                    num = nxt[before]
                    before = num

            prev[nxt[num]], prev[front[dst]] = 0, num
            front[src], front[dst], nxt[num] = nxt[num], front[src], front[dst]

            len_conv[src] -= idx
            len_conv[dst] += idx

            if len_conv[dst] == idx:
                back[dst] = num

        print(len_conv[dst])

    elif order[0] == 500:
        num = order[1]
        a = prev[num] if prev[num] != 0 else -1
        b = nxt[num] if nxt[num] != 0 else -1
        print(a + 2*b)

    else:
        num = order[1]
        if len_conv[num] == 0:
            a, b, c = -1, -1, 0
        else:
            a, b, c = front[num], back[num], len_conv[num]
        print(a + 2*b + 3*c)