[코드트리/파이썬] 가로등 설치

https://www.codetree.ai/ko/frequent-problems/samsung-sw/problems/street-light-installation/description

코딩테스트 기출 문제 설명: 가로등 설치 | 코드트리

 

 

import heapq
def init(query : list[int]) :
    """
    cmd = 100 일때 초기 세팅
    """
    cmd, N, M, *town = query
    town = [-1] + town      # [-1, n1, n2, n3, ...] idx : 가로등 번호, n_ : 가로등 위치
    distances = []
    heapq.heappush(distances, (-abs(town[1])*2, 0, 1))    # 거리의 최댓값, left, right (가로등 번호)
    heapq.heappush(distances, (-abs(town[-1] - N)*2, M, 0))
    for i in range(1, len(town)-1) :
        heapq.heappush(distances, (-abs(town[i+1] - town[i]), i, i+1))

    return N, M, town, distances

def add_lamp(distances: list[int], town: list[int], dist_dict: dict) :
    """
    cmd = 200 일때, 램프 추가
    """
    dist, left, right = 0, 0, 0
    for i in range(len(distances)) :
        dist, left, right = distances[i] # 힙으로 넣었으면 가장 작은 값(거리는 음수니까 절대값은 가장 큰 값)이 나온다.
        if left != 0 and right != 0 and dist_dict[distances[i]] != -1:
            break
    dist = -dist
    
    town.append(town[left] + dist//2)   # 고유번호 : len(town) - 1, value : lamp 위치 
    curr_lamp_idx = len(town)-1
    heapq.heappush(distances, (-abs(town[left] - town[curr_lamp_idx]), left, curr_lamp_idx))
    heapq.heappush(distances, (-abs(town[right] - town[curr_lamp_idx]), curr_lamp_idx, right))
    
    dist_dict[distances[i]] = -1
    dist_dict[(-abs(town[left] - town[curr_lamp_idx]), left, curr_lamp_idx)] = 1
    dist_dict[(-abs(town[right] - town[curr_lamp_idx]), curr_lamp_idx, right)] = 1

    return distances, town, dist_dict

def delete_lamp(distances: list[int], town: list[int], dist_dict: dict, D:int, N:int) :
    """
    cmd = 300일 때, 램프 삭제
    D : 삭제할 가로등 인덱스
    N : 마을 크기
    """
    town[D] = -1
    temp_left = temp_right = 0  # 뽑힌거 제거하고 남아있는 가로등 인덱스

    for i in range(len(distances)) :
        dist, left, right = distances[i]
        if dist_dict[distances[i]] != -1 :
            if left == D :   # 2개가 삭제 될거다.
                dist_dict[distances[i]] = -1    # 삭제
                temp_right = right

            elif right == D :
                dist_dict[distances[i]] = -1    # 삭제
                temp_left = left

    if temp_left == 0 :
        heapq.heappush(distances, (-abs(town[temp_right]) * 2, temp_left, temp_right))
        dist_dict[(-abs(town[temp_right]) * 2, temp_left, temp_right)] = 1
    elif temp_right == 0 :
        heapq.heappush(distances, (-abs(town[temp_left]-N) * 2, temp_left, temp_right))
        dist_dict[(-abs(town[temp_left]-N) * 2, temp_left, temp_right)] = 1
    else :
        heapq.heappush(distances, (-abs(town[temp_left]-town[temp_right]), temp_left, temp_right))
        dist_dict[(-abs(town[temp_left]-town[temp_right]), temp_left, temp_right)] = 1
    
    return distances, town, dist_dict

def get_distances(distances: list[int], town: list[int], dist_dict: dict) :
    while distances:
        dist, left, right = distances[0]
        if dist_dict[distances[0]] != -1 :
            return -dist
        
        heapq.heappop(distances)

# 거리 계산은 
def main() :
    Q = int(input())    # 쿼리 개수
    # 가로등 간의 거리를 저장
    N = M = 0       # N, M 
    distances = []  # 거리 저장을 어떤식으로 해야할까 (-거리, left, right) 이런식으로 해야되나
    town = []       # 가로등 위치 정보
    dist_dict = {}  # 이 거리가 유효한건지 아닌지 검사하기 위한 dict
    for _ in range(Q) :
        query = list(map(int, input().split()))
        if query[0] == 100 :
            N, M, town, distances = init(query)
            for i in range(len(distances)) :
                dist_dict[distances[i]] = 1

        elif query[0] == 200 :
            distances, town, dist_dict = add_lamp(distances, town, dist_dict)

        elif query[0] == 300 :
            cmd, D = query
            # pop 하는건 인덱스도 다시 초기화 되잖아
            distances, town, dist_dict = delete_lamp(distances, town, dist_dict, D, N)

        elif query[0] == 400 :
            print(get_distances(distances, town, dist_dict))
            
    
if __name__ == "__main__" :
    main()

이 문제를 4시간동안 고민하고 풀었는데 주어진 테스트 케이스만 맞고 돌리니까 바로 틀렸다는 답안이 보여졌다. 정답처리하는 테스트케이스를 보고나니 뭔가 잘못 작성한 코드라는 것이 느껴졌다. 이따 스터디 하면서 다듬어봐야겠다.