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

문제 링크

https://www.codetree.ai/training-field/frequent-problems/problems/santa-gift-factory-2/description?page=1&pageSize=20 

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

 

풀이

계속 테스트케이스 4번에서 틀려서 디버깅하는데 엄청 오래걸림. 

4번은 틀리는 조건이 하나만 있는 게 아니라, 하나 수정하고 또 수정하고 반복 ,, 

질문에 올라와 있는 여러 TC 돌려가며 디버깅 했다. 

박스가 하나 밖에 안 남아있는 경우 or 빈 벨트에 박스가 추가되는 경우에 tail도 갱신해줘야 한다. head만 해서 틀린 경우가 많았다. 오열 

시간 나면 코드 다시 수정해야징 

 

pre : id에 해당하는 상자 별 pre 상자 id 저장

nxt : id에 해당하는 상자 별 next 상자 id 저장

head : 벨트 별 head 상자 id 저장 

tail : 벨트 별 tail 상자 id 저장

num_gift : 벨트 별 상자 개수 저장 

pre = defaultdict(int)
nxt = defaultdict(int)

head = [0] * (MAX_N+1)
tail = [0] * (MAX_N+1)
num_gift = [0] * (MAX_N+1)

 

1. 공장 설립 

def build(data):
    global n, m
    
    n = data[1]
    m = data[2]
    
    nums = data[3:]
    
    for i in range(1,m+1):
        b_num = nums[i-1]
        num_gift[b_num] += 1

        # 이미 존재하면 뒤에 연결 
        if tail[b_num] :
            t_id = tail[b_num]
            tail[b_num] = i
            nxt[t_id] = i
            pre[i] = t_id 
        # 벨트에 아무 박스도 존재하지 않는 경우 
        else :
            tail[b_num] = i
            head[b_num] = i

2. 물건 모두 옮기기 

def move_all(data):
    m_src = data[1]
    m_dst = data[2]
    
    # 옮길 물건 존재하는 경우 
    if head[m_src] != 0 : 
        src_tail = tail[m_src]
        dst_head = head[m_dst]
        
        # head 갱신 
        head[m_dst] = head[m_src]
        # 옮긴 위치에 박스가 하나도 없는 경우에 tail도 갱신해줘야 한다
        if tail[m_dst] == 0 :
            tail[m_dst] = tail[m_src]
            
        # pre,nxt 갱신 
        pre[dst_head] = src_tail
        nxt[src_tail] = dst_head 
        
        # 기존 벨트 head, tail 삭제 
        head[m_src], tail[m_src] = 0,0 
        
        # 벨트 개수 갱신 
        num_gift[m_dst] += num_gift[m_src]
        num_gift[m_src] = 0

    print(num_gift[m_dst])

3. 앞 물건만 교체하기 

# 서로 교환할 때 함수 
def change_head(m_src, m_dst):
    src_head = head[m_src]
    dst_head = head[m_dst]
    
    # 하나만 있었던 경우 tail 갱신 
    if nxt[src_head] == 0 :
        tail[m_src] = dst_head
    else :
        pre[nxt[src_head]] = dst_head 
    
    # 하나만 있었던 경우 tail 갱신 
    if nxt[dst_head] == 0 :
        tail[m_dst] = src_head
    else :
        pre[nxt[dst_head]] = src_head
    
    src_nxt = nxt[src_head]
    dst_nxt = nxt[dst_head]

    nxt[dst_head] = src_nxt 
    nxt[src_head] = dst_nxt
        
    head[m_src] = dst_head 
    head[m_dst] = src_head 
    
def change(data):
    m_src = data[1]
    m_dst = data[2]
	
    # 둘다 빈 벨트인 경우 옮기지 않음 
    if head[m_src] == 0 and head[m_dst] == 0 :
        print(num_gift[m_dst])
        return 
	
    # 하나 빈 벨트인 경우 
    if head[m_src] == 0 :
        dst_head = head[m_dst]
        
        head[m_dst] = nxt[dst_head]
        head[m_src] = dst_head
        tail[m_src] = dst_head 
        
        # 물건이 하나 밖에 없는 경우 tail도 갱신 
        if nxt[dst_head] == 0:
            tail[m_dst] = 0
        else :        
            pre[nxt[dst_head]] = 0 
        
        nxt[dst_head] = 0
        
        num_gift[m_src] += 1
        num_gift[m_dst] -= 1 
    # 하나 빈 벨트인 경우 
    elif head[m_dst] == 0 :
        src_head = head[m_src]
        
        head[m_src] = nxt[src_head]
        head[m_dst] = src_head
        tail[m_dst] = src_head
        
        # 물건이 하나 밖에 없는 경우 tail도 갱신 
        if nxt[src_head] == 0 :
            tail[m_src] = 0 
        else:
            pre[nxt[src_head]] = 0 
        
        nxt[src_head] = 0
        
        num_gift[m_src] -= 1
        num_gift[m_dst] += 1
    # 서로 교환하는 경우 
    else :
        change_head(m_src,m_dst)
    
    print(num_gift[m_dst])

4. 물건 나누기

def divide_gift(data):
    m_src = data[1]
    m_dst = data[2]
    
    n = num_gift[m_src]
    
    # 상자 옮기지 않음 
    if n <= 1 :
        print(num_gift[m_dst])
        return 
    
    cnt = n // 2 
    # 옮겨야 하는 마지막 번호 찾기 
    last_id = head[m_src]
    
    for _ in range(cnt-1):
        last_id = nxt[last_id]
 	
    src_head = head[m_src]
    dst_head = head[m_dst]
	
    # head 갱신 
    head[m_dst] = src_head 
    head[m_src] = nxt[last_id]
    
    # 빈 벨트로 이동한 경우 tail 갱신 
    if tail[m_dst] == 0 :
        tail[m_dst] = last_id
    
    # 빈 벨트가 된 경우 tail 갱신 
    if nxt[last_id] == 0:
        tail[m_src] = 0 
    else :
        pre[nxt[last_id]] = 0 
        
    pre[dst_head] = last_id
    nxt[last_id] = dst_head
    
    num_gift[m_dst] += cnt 
    num_gift[m_src] -= cnt 
    
    print(num_gift[m_dst])

5. 선물 정보 얻기

def get_gift(data):
    p_num = data[1]
    
    a = -1
    b = -1 
    if pre[p_num] :
        a = pre[p_num]
    if nxt[p_num]:
        b = nxt[p_num]
        
    print(a+2*b)

6. 벨트 정보 얻기 

def get_belt(data):
    b_num = data[1]
    
    if num_gift[b_num] == 0:
        a = -1 
        b = -1 
        print(a+2*b)
        return 
    
    a = head[b_num]
    b = tail[b_num]
    c = num_gift[b_num]
    
    print(a+2*b+3*c)

 

전체 코드는 아래.. 

from collections import defaultdict 

MAX_N = 10
MAX_M = 10

q = int(input())

n, m = -1,-1

pre = defaultdict(int)
nxt = defaultdict(int)

head = [0] * (MAX_N+1)
tail = [0] * (MAX_N+1)
num_gift = [0] * (MAX_N+1)

def build(data):
    global n, m
    
    n = data[1]
    m = data[2]
    
    nums = data[3:]
    
    for i in range(1,m+1):
        b_num = nums[i-1]
        num_gift[b_num] += 1

        # 이미 존재하면 뒤에 연결 
        if tail[b_num] :
            t_id = tail[b_num]
            tail[b_num] = i
            nxt[t_id] = i
            pre[i] = t_id 
        # 아무것도 존재하지 않음 
        else :
            tail[b_num] = i
            head[b_num] = i 
    
def move_all(data):
    m_src = data[1]
    m_dst = data[2]
    
    # 물건 존재
    if head[m_src] != 0 : 
        src_tail = tail[m_src]
        dst_head = head[m_dst]
        
        head[m_dst] = head[m_src]
        if tail[m_dst] == 0 :
            tail[m_dst] = tail[m_src]
            
        pre[dst_head] = src_tail
        nxt[src_tail] = dst_head 
        
        head[m_src], tail[m_src] = 0,0 
        
        num_gift[m_dst] += num_gift[m_src]
        num_gift[m_src] = 0

    print(num_gift[m_dst])
    
def change_head(m_src, m_dst):
    src_head = head[m_src]
    dst_head = head[m_dst]
    
    # 하나만 있었던 경우 
    if nxt[src_head] == 0 :
        tail[m_src] = dst_head
    else :
        pre[nxt[src_head]] = dst_head 
    
    # 하나만 있었던 경우 
    if nxt[dst_head] == 0 :
        tail[m_dst] = src_head
    else :
        pre[nxt[dst_head]] = src_head
    
    src_nxt = nxt[src_head]
    dst_nxt = nxt[dst_head]

    nxt[dst_head] = src_nxt 
    nxt[src_head] = dst_nxt
        
    head[m_src] = dst_head 
    head[m_dst] = src_head 
    
def change(data):
    m_src = data[1]
    m_dst = data[2]

    if head[m_src] == 0 and head[m_dst] == 0 :
        print(num_gift[m_dst])
        return 

    if head[m_src] == 0 :
        dst_head = head[m_dst]
        
        head[m_dst] = nxt[dst_head]
        head[m_src] = dst_head
        tail[m_src] = dst_head 
        
        # 물건이 하나 밖에 없는 경우 
        if nxt[dst_head] == 0:
            tail[m_dst] = 0
        else :        
            pre[nxt[dst_head]] = 0 
        
        nxt[dst_head] = 0
        
        num_gift[m_src] += 1
        num_gift[m_dst] -= 1 
        
    elif head[m_dst] == 0 :
        src_head = head[m_src]
        
        head[m_src] = nxt[src_head]
        head[m_dst] = src_head
        tail[m_dst] = src_head
        
        if nxt[src_head] == 0 :
            tail[m_src] = 0 
        else:
            pre[nxt[src_head]] = 0 
        
        nxt[src_head] = 0
        
        num_gift[m_src] -= 1
        num_gift[m_dst] += 1
    else :
        change_head(m_src,m_dst)
    
    print(num_gift[m_dst])

def divide_gift(data):
    m_src = data[1]
    m_dst = data[2]
    
    n = num_gift[m_src]
    
    if n <= 1 :
        print(num_gift[m_dst])
        return 
    
    cnt = n // 2 
    # 옮겨야 하는 마지막 번호
    last_id = head[m_src]
    
    for _ in range(cnt-1):
        last_id = nxt[last_id]
    
    print(m_src, m_dst, last_id)
    src_head = head[m_src]
    dst_head = head[m_dst]

    head[m_dst] = src_head 
    head[m_src] = nxt[last_id]
    
    if tail[m_dst] == 0 :
        tail[m_dst] = last_id
        
    if nxt[last_id] == 0:
        tail[m_src] = 0 
    else :
        pre[nxt[last_id]] = 0 
        
    pre[dst_head] = last_id
    nxt[last_id] = dst_head
    
    num_gift[m_dst] += cnt 
    num_gift[m_src] -= cnt 
    
    print(num_gift[m_dst])
    
def get_gift(data):
    p_num = data[1]
    
    a = -1
    b = -1 
    if pre[p_num] :
        a = pre[p_num]
    if nxt[p_num]:
        b = nxt[p_num]
        
    print(a+2*b)
    
def get_belt(data):
    b_num = data[1]
    
    if num_gift[b_num] == 0:
        a = -1 
        b = -1 
        print(a+2*b)
        return 
    
    a = head[b_num]
    b = tail[b_num]
    c = num_gift[b_num]
    
    print(a+2*b+3*c)
    
for _ in range(q):
    data = list(map(int,input().split()))
    q_type = data[0]
    
    if q_type == 100:
        build(data)
    if q_type == 200 :
        move_all(data)
    if q_type == 300 :
        change(data)
    if q_type == 400:
        divide_gift(data)
    if q_type == 500:
        get_gift(data)
    if q_type == 600:
        get_belt(data)