2513. 통학버스 (Python)

2513번: 통학버스

 

풀이

 학교의 위치부터 가장 멀리 떨어져 있는 아파트부터 먼저 방문해서 태울 수 있는 만큼 학생들을 태워가며 버스의 이동거리를 계산해주면 된다. 이때 학교를 기준으로 왼쪽에 있는 아파트와 오른쪽에 있는 아파트로 나누어 왼쪽으로 한 번 작업을 수행하고 오른쪽으로 한 번 작업을 수행해주면 된다. 

 학교로부터 가장 멀리 떨어져 있는 아파트부터 먼저 방문해야 하므로 왼쪽 아파트의 경우에는 아파트 위치를 기준으로 최소힙으로 구성하였고 오른쪽 아파트의 경우에는 최대힙으로 구성하였다. 힙에서 가장 위에 있는 아파트 정보를 가져와 학생들을 태우고 해당 아파트의 모든 학생들을 태웠다면 그대로 힙에서 제거해주고, 아직 남은 학생들이 있다면 다시 힙에 아파트 정보를 넣어주는 방식으로 체크하면 된다.

 

코드

import sys
import heapq
input = sys.stdin.readline


# students: 학생들이 사는 아파트에 대한 정보를 담아 놓은 힙 
# base: 최대 힙일 경우에는 -1을 곱해주는 작업이 필요하기 때문에 사용함. 최소 힙일때는 1이 들어오고 최대 힙일때는 -1이 들어온다.
def goSchool(students, base):
    global res
    nowInBus = K  # 현재 버스에 태울 수 있는 학생의 수

    while students:
        student = heapq.heappop(students)
        apartment, studentNum = student[0] * base, student[1] # 최대힙일 경우 아파트 위치가 음수이기 때문에 -1을 곱해준다.

        # 남은 정원이 버스의 정원과 일치한다는 것은 버스가 학교부터 새롭게 출발하여 학생들을 태우기 시작했다는 것을 의미한다. 
        # 그렇기 때문에 이동 거리를 현재 아파트 위치를 기준으로 더해준다. 
        if nowInBus == K: 
            res += abs(S - apartment) * 2

        # 현재 아파트의 학생 수가 버스에 태울 수 있는 정원 이하일 경우에는 
        # 모든 학생을 버스에 태운다.  
        if studentNum <= nowInBus:
            nowInBus -= studentNum
            studentNum = 0 # 모든 학생이 탔기 때문에 현재 아파트에 남은 학생의 수는 없다. 

            if nowInBus == 0: # 만약 버스가 꽉 찼다면 학교에 학생을 내려주고 다시 버스의 정원을 K로 늘려준다.
                nowInBus = K
        else: # 현재 아파트의 학생 수가 현재 버스에 태울 수 있는 정원을 초과할 경우 
            studentNum -= nowInBus # 남은 정원만큼만 학생을 태우고
            nowInBus = K # 버스는 꽉 찼으므로 학생들을 내려주고 다시 학교에서 출발하여 학생들을 태운다. 

        if studentNum != 0: # 현재 아파트에 아직 태울 수 있는 학생이 남았다면 다시 힙에 넣어준다. 
            heapq.heappush(students, (base * apartment, studentNum))


if __name__ == '__main__':
    # 아파트 단지의 수, 버스의 정원, 학교의 위치
    N, K, S = map(int, input().split())
    leftStudents = [] # 학교 위치를 기준으로 왼쪽에 사는 학생들 정보. 위치에 대하여 최소 힙으로 구성
    rightStudents = [] # 학교 위치를 기준으로 오른쪽에 사는 학생들 정보. 위치에 대하여 최대 힙으로 구성

    for _ in range(N):
        a, b = map(int, input().split())

        if a < S:
            heapq.heappush(leftStudents, (a, b))
        else:
            heapq.heappush(rightStudents, (-a, b))

    res = 0
    goSchool(leftStudents, 1) # 학교를 기준으로 왼쪽 학생들을 태우는 작업을 한다.
    goSchool(rightStudents, -1) # 학교를 기준으로 오른쪽 학생들을 태우는 작업을 한다. 

    print(res)