[Python] 백준 29767. 점수를 최대로

문제

https://www.acmicpc.net/problem/29767

단대소프트고에는 교실 N개가 있다. 교실은 1번부터 N번까지 1,2,…,N 순서로 연달아 있다.

학교 밖에는 K명의 학생들이 있다. K명의 학생은 학교에 들어가기 전 학생마다 목적지 교실을 정하게 된다. j번째 학생의 목적지는 Bj번 교실이다. 학생들의 목적지는 모두 달라야 한다. j번째 학생은 1번 교실, 2번 교실, …,Bj번 교실까지 모든 교실에 들어갔다 나오게 된다. 마지막 교실까지 들어갔다 나오게 되면 아무 교실도 방문하지 않고 다시 학교 밖으로 나간다.

모든 학생의 방문이 끝나면, 교실의 점수를 구할 수 있다. i번째 교실의 점수는 Ai × (i번째 교실에 학생이 들어갔다 나온 횟수)가 된다. 학생들은 학교에 들어가기 전 방문 후의 교실의 점수의 합이 최대가 되도록 의논하여 목적지를 정한다. 이때 방문이 끝난 후 모든 교실의 점수의 합을 구해보자.

입력

첫째 줄에 N과 K가 공백으로 구분되어 입력된다. (1≤K≤N≤300000)

둘째 줄에 정수 A1,A2,...AN이 공백으로 구분되어 입력된다. (−108≤Ai≤108)

출력

첫째 줄에 방문이 끝난 후 모든 교실의 점수의 합을 출력한다.

---

풀이

처음에는 조금 헤매서 알고리즘 분류를 보니 누적합이 적혀있었다. 누적합... 누적합이 왜 필요할까...어떤 기준으로 합이 더 큰걸 고르지 고민을 하다가 일단 예제로 직접 시도를 해봤다. 

각각의 위치에서 최선을 다했다고 해보자.

1) 그러면 1번째 학생은 4까지 모든 교실을 다 들를 것이고,

2) 2번째 학생은 굳이 -1까지 갈 필요가 없기 때문에 2에서 멈출 것이다.

3) 그리고 3번째 학생도 같은 이유로 3에서 멈출 것이다.

 

각 학생들이 들린 교실의 a[i]값을 더해보면... 

k = 1 : 3 / k = 2 : 0 / k = 3 : 3 ===> 으로 합은 6이다. 그러니까 정리하자면 i번째까지의 누적합을 내림차순으로 나열한 다음에 k번까지만 더해주게 되면 방문 후의 교실의 합이 최대가 되는 값을 구할 수 있다는 것이다. 

 

이해가 안될 수 있어서 추가적으로 더 적자면, 여기서 누적합은 각각의 교실에 도달했을때 지나간 교실 A 값들의 합이기 때문에, 문제에서 주어진 Ai * (i번째 교실에 학생이 들어간 횟수)랑 동일하다. 그리고 내림차순으로 k번째까지 더하는 건 위 예시의 누적합을 배열로 표현하면 이해하기 쉬울 것 같다.

[ 3, -2, 0, -1, 3 ] 일 것이기 때문에 내림차순으로 높은 값 k까지만 더해줘서 최대 값을 구하는 것이다.

 

import sys
input = sys.stdin.readline

n, k = map(int, input().split())
a = list(map(int, input().split()))
result = 0

nums = [0]*n
nums[0] = a[0]
for i in range(1, n):
    nums[i] = nums[i-1] + a[i]

nums.sort(reverse=True)

for i in range(k):
    result += nums[i]

print(result)

 

시간복잡도

두 for문은 각각 O(n), O(k)이지만 sort 정렬 시간복잡도가 O(nlogn)이기 때문에 O(nlogn)일 것 같다.