지형 편집

import itertools from math import inf from collections import Counter def solution(land, P, Q): N = len(land) # 우선 땅을 한번 스캔해서, # n층에 땅이 몇 칸인지 세는 카운터를 만듭니다. # 10^9개 층에 대해 저장할 수는 없지만, # "실제로 등장하는" 층에 대해서만 저장해두면 충분합니다. cnt = Counter(itertools.chain(*land)) levels = list(sorted(cnt.keys())) total = sum(lvl * n for lvl, n in cnt.items()) levelcnt = {} # levelcnt[n] : n층에 땅이 몇 칸인지? cumsum = 0 for level in reversed(levels): cumsum += cnt[level] levelcnt[level] = cumsum # 0층으로 만들때의 cost부터 계산해 봅시다. # 결국 모든 블록을 제거해야 하므로, # cost는 Q * total 입니다. cost = Q * total answer = inf prev_level = 0 # 이제 가장 낮은 층 부터 시작해서, levels 배열을 오름차순으로 봅니다. for level in levels: # 이전 level과 비교해봤을 때, 이 level에서 제거해야 하는 블록 수가 # 얼마나 줄어들었을지 계산해봅니다. # 총 (level - prev_level) * levelcnt[level] 개입니다. cost -= Q * (level - prev_level) * levelcnt[level] # 그 대신 더 쌓아야 하는 블록 수가 얼마나 늘어났을지 계산해봅니다. # 총 (level - prev_level) * (N**2 - levelcnt[level]) 개입니다. cost += P * (level - prev_level) * (N**2 - levelcnt[level]) answer = min(answer, cost) prev_level = level return answer