[Python/파이썬] 백준 2900번 프로그램

2900번: 프로그램

 

문제

창영이가 에러를 찾기 위해서 디버깅을 하고 있다. 이 프로그램은 크기가 N이고 0으로 채워져있는 배열을 a를 만들고, 아래 something 함수를 호출한다.

void something(int jump) {
    int i = 0;
    while (i < N) {
        a[i] = a[i] + 1;
        i = i + jump;
    }
}

창영이는 함수를 K번 호출하려고 한다. 각각 호출할 때, 인자로 넘기는 jump의 값은 $X_1, X_2, X_3, ... X_k$ 순서이다.

이렇게 호출한 뒤에는 배열의 값이 정상적으로 들어갔는지를 확인하려고 한다. 확인은 총 Q번 하고, 매번 확인을 할 때마다 L과 R(L ≤ R)을 정한뒤, 그 구간의 배열의 합을 구한다. 즉, $a[L] + a[L+1] + ... + a[R]$을 구한다.

함수를 호출할 때 필요한 X의 값과 창영이가 확인한 횟수 Q가 주어졌을 때, 확인한 결과(배열의 합)을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 배열의 크기 N과 함수의 호출 횟수 K가 주어진다. $(1 ≤ N, K ≤ 10_6)$

둘째 줄에 함수를 호출할 때 사용하는 인자의 값 $X_1, X_2, X_3, ... X_k$가 공백으로 구분되어 주어진다. $(1 ≤ X_i < N)$

셋째 줄에는 확인하는 횟수 Q가 주어진다. $(1 ≤ Q ≤ 10_6)$

넷째 줄부터 Q개 줄에는 각 확인에 사용하는 L과 R이 주어진다. (0 ≤ L ≤ R < N)

출력

출력은 총 Q줄이다. 창영이가 확인하는데 사용한 L과 R이 주어졌을 때, $a[L] + a[L+1] + a[L+2] ... + a[R]$을 출력한다. 

코드

from collections import defaultdict
import sys
input = sys.stdin.readline

n, k = map(int, input().split())
x_lst = list(map(int, input().split()))
x_dict = defaultdict(int)
for x in x_lst:
    x_dict[x] += 1
# something 함수
a = [0]*n
for jumping, num in x_dict.items():
    for i in range(0, n, jumping):
        a[i] += num
# 누적합
ps = [0] * (n+1)
for i in range(1, n+1):
    ps[i] = ps[i-1]+a[i-1]
q = int(input())

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

 

N, K의 범위가 $(1 ≤ N, K ≤ 10_6)$이므로 $X_1, X_2, X_3, ... X_k$을 모두 한 번씩 something 함수의 인자로 넣고 실행한다면 시간 초과가 날 것이 분명하다. 그러므로 X 배열 내의 같은 수의 개수를 세서 같은 수에 대해서는 something을 한 번만 실행하도록 한다. 이렇게 한다면 a[i] = a[i]+1에서 1 대신에 같은 수의 개수를 더해주어야 한다.

이렇게 something을 모두 실행한 결과가 a에 저장되면 a의 누적합을 미리 구해서 L 이상 R 이하의 배열의 합을 구할 때 사용하도록 한다.