Koala - 6기/코딩테스트 준비 스터디

[백준 / python] 17390번: 이건 꼭 풀어야 해!

Juno7 2022. 5. 15. 12:39

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

 

17390번: 이건 꼭 풀어야 해!

[2, 5, 1, 2, 3]을 비내림차순으로 정렬하면 [1, 2, 2, 3, 5]이다.

www.acmicpc.net

 

문제분석

 

수열을 입력받고 이를 비 내림차순으로 정렬해서 수열을 만든다. 그런 다음 정수를 두 개 입력받는다. 수열에서 입력받은 정수 인덱스에 해당하는 값들의 누적 합을 구해서 출력하면 된다.

 

코드

from itertools import accumulate
input = __import__('sys').stdin.readline

n, q = map(int, input().split())
a = sorted([*map(int, input().split())])
p_sum = list(accumulate(a))
p_sum.insert(0, 0)

for _ in range(q):
    l, r = map(int, input().split())
    l -= 1
    print(p_sum[r] - p_sum[l])

 

 

문제풀이

 

누적합을 구해주는 itertools의 accumulate를 활용한다. n, q를 입력받아 함수의 길이와 케이스 크기를 입력받는다. 그런 다음 수열을 입력받고 이에 대한 누적합을 구하기 위해 p_sum에 accumulate를 활용한 누적합 리스트를 저장한다. 두 개의 누적합을 빼서 답을 도출해야 하는데 만약 처음부터 시작하는 누적합을 구할 경우, 의도치 않게 첫 번째 인덱스를 빼게 된다. 이를 방지하기 위해 누적합 수열에 맨 처음에는 0을 넣어준다. 그다음 loop를 돌면서 l인덱스에 1을 뺀 값과 r인덱스를 입력받아 누적합을 구한다. 예를 들어, 2, 4를 입력받을 경우 네 번째까지의 누적 합에서 첫 번째까지의 누적합을 뺀 값을 출력하기 때문에 구하려는 값이 나오게 할 수 있다.