BOJ 26123 - 외계 침략자 윤이 (Python3)

이 문제는 UNIST 대회인 Uni-CODE 2022의 B번 문제이다. 대회 때 완벽하게 풀어 낸 유일한 문제이다.

A번은 어떻게 푸는지 감도 못 잡겠더라. 정해가 빨리 공개되면 좋겠다.

 

그럼 본격적으로 풀이에 돌입해 보자.

---

문제

외계인 윤이는 지구를 정복하고자 세계의 중심 도시인 울산을 침략했다. 울산에는 N개의 빌딩이 일렬로 늘어서 있고, 왼쪽에서 i번째 건물의 높이는 \(h_i\)이다.

윤이는 울산을 파괴하기 위해 다음과 같은 계획을 세웠다. 윤이는 매일 UFO를 타고 울산의 상공을 가르며 가장 높이가 높은 빌딩에 레이저를 발사할 것이다. 레이저에 맞은 빌딩은 높이가 1 낮아진다. 만약 그 날에 가장 높이가 높은 빌딩이 여러 개라면, 해당하는 모든 빌딩에 레이저를 발사한다. 만약 이미 모든 빌딩의 높이가 0이 되었다면, 윤이는 더 이상 레이저를 발사하지 않는다.

윤이는 지금까지 D일 동안 계획을 착실하게 수행해 왔다. 윤이가 D일 동안 레이저를 발사한 총 횟수를 구하시오.

 

입력

첫 번째 줄에 건물의 개수 N과 윤이가 계획을 수행한 날의 수 D가 공백으로 구분되어 주어진다. (1 ≤ N ≤ 300000, 1 ≤ D ≤ 300000)

두 번째 줄에 건물의 높이를 나타내는 N개의 정수 \(h_i\)가 공백으로 구분되어 주어진다. (1 ≤ \(h_i\) ≤ 300000)

 

예제 입력)

# 예제 입력 1
5 3
1 3 2 5 4

# 예제 입력 2
5 2
1 1 1 1 1

 

출력

윤이가 D일 동안 레이저를 발사한 총 횟수를 출력한다. (32비트 정수 범위를 넘을 수 있음에 유의하라.)

 

예제 출력)

# 예제 출력 1
6

# 예제 출력 2
5

예제 1의 경우 처음 건물의 높이는 [1,3,2,5,4]이다. 첫 번째 날에 윤이는 레이저를 1회 발사하여 건물의 높이가 [1,3,2,4,4]가 되며, 두 번째 날에 윤이는 레이저를 2회 발사하며 건물의 높이가 [1,3,2,3,3]이 된다. 세 번째 날에 윤이는 레이저를 3회 발사하며, 건물의 높이는 [1,2,2,2,2]가 된다.

예제 2의 경우 처음 건물의 높이는 [1,1,1,1,1]이다. 첫 번째 날에 윤이는 레이저를 5회 발사하여 건물의 높이는 [0,0,0,0,0]이 된다. 두 번째 날에 이미 모든 건물의 높이가 0이므로 윤이는 레이저를 발사하지 않는다.

 

---

내 코드

dic = {}
n,d = map(int,input().split())
l,ans = list(map(int,input().split())),0
for i in l:
    try: dic[i] += 1
    except KeyError: dic[i] = 1
x = max(list(dic.keys()))
for _ in range(d):
    if x == 0: print(ans); quit()
    ans += dic[x]
    try: dic[x-1] += dic[x]
    except KeyError: dic[x-1] = dic[x]
    x -= 1
print(ans)

 

이번에 사용할 것은 딕셔너리. 높이가 key인 건물의 수를 value로, dic = {key: value, ...} 식으로 저장하게 될 것이다.

첫 번째 반복문에 건물의 높이인 리스트 l을 순회할 것이다. l의 원소 i는 높이가 i인 건물을 나타내는 것이니, 높이가 key인 건물의 수를 dic[key]로 저장하는 이 코드의 특성 상 dic[i]에 1을 추가해야 한다.

만약 그럴 수 없을 때는 높이가 i인 건물이 아직 없다는 의미이므로 try-except문으로 예외처리하여 그 값을 1로 만들어 주면 끝.

 

그 다음 for문에서는 가장 높은 건물의 갯수를 답에 d번동안 더해줄 것이다.

x는 가장 높은 건물의 높이이고, 이게 한 번 for문을 돌면 1이 줄어들어야만 한다. 윤이가 레이저를 날려서 가장 윗층을 날려버렸을 테니.

그러면 그 건물들은 모두 x-1의 높이를 가지게 되므로, dic[x-1]에 dic[x]를 더해 주어야 한다.

만약 x-1이 dic의 key가 아니라면 마찬가지로 건물의 높이가 x-1인 것들이 없었다는 것을 의미하니, 예외처리하여 그 값을 dic[x]로 만들어 주면 끝.

x=0일 때 프로그램을 끝내는 것을 잊지 말자!

 

이로서 26123번의 풀이를 마친다.

그럼, 오늘도 당신의 코딩 실력이 상승하기를.