누적합

길이가 N인 리스트에서 누적합을 구해서 리턴하는 함수를 구하세요.
1부터 N까지 숫자로 이뤄진 길이가 N인 리스트와, 구하려는 누적합의 시작하는 지점과 끝나는 지점을 담은 M개의 쿼리 데이터를 받습니다.

1
2
3
4
N = 5
M = 5
data = [10, 20, 30, 40, 50]
queries = [[1, 3], [2, 4], [3, 5], [1, 5], [4, 4]]

일때 결과는

data에서

  • 1번째부터 3번째까지의 누적합 : 60
  • 2번째부터 4번째까지의 누적합 : 90
  • 3번째부터 5번째까지의 누적합 : 120
  • 1번째부터 5번째까지의 누적합 : 150
  • 4번째부터 4번째까지의 누적합 : 40

[60, 90, 120, 150, 40]을 리턴합니다.

조건

def solution(data, queries):

N은 data의 길이, data요소의 최대값

1<=N<=십만

M은 queries의 길이

1<=M<=만

나의 코드

1
2
3
4
N = 5
M = 5
data = [10, 20, 30, 40, 50]
queries = [[1, 3], [2, 4], [3, 5], [1, 5], [4, 4]]
1
2
3
4
5
6
def solution(data, queries):
ls = []
for i,j in queries:
ls.append(sum(data[i-1:j]))

return ls

주어진 조건 중 최대 길이의 sample 생성

1
2
3
4
5
6
from random import randint

N = 100000
M = 10000
data = [randint(1, N) for _ in range(N)]
q = [ sorted([randint(1, N), randint(1, N)]) for _ in range(M) ]

실행시간에서 S가 나오면 의심해 봐야 한다.

효율적인 코드

prefix_sum: 사전에 미리 누적합을 구해놓는다.

내가 다시 푼 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
def solution(data, queries):
ls = [0]
total = 0
for num in data:
total += num
ls.append(total)

result = []

for start, end in queries:
result.append(ls[end] - ls[start-1])

return result

아래와 같이 누적합을 itertools 패키지를 통해 구할 수 있다.

1
2
3
4
from itertools import accumulate

acc_data = [0] + list(accumulate(data))
acc_data
Share