[BOJ] 10986 나머지 합 (Python3)

문제

수 N개 A1, A2, …, AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, Ai + … + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10^6, 2 ≤ M ≤ 10^3)

둘째 줄에 N개의 수 A1, A2, …, AN이 주어진다. (0 ≤ Ai ≤ 10^9)

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

첫번째 풀이 코드

n, m = map(int, input().split())
numbers = list(map(int, input().split()))

answer = 0

dp = [0 for i in range(n + 1)]
dp[1] = numbers[0]

for i in range(1, n):
    number = numbers[i]
    for j in reversed(range(i + 1)):
        value = dp[j] + number
        
        if value % 3 == 0:
            answer += 1
            
        dp[j + 1] = value
        
print(answer)

 5 3
 1 2 3 1 2


7

첫번째 풀이과정

얼마전에 이런 구간에 대한 값을 찾는 비슷한 문제를 dp로 푼 기억이 있어서, 처음에는 위와같이 dp로 풀었다. 이전 숫자까지의 합에 드는 과정을 기록하며 제거하면 통과할 수 있다고 생각했는데, 계속 시간초과가 났다.

그래서 다시 생각해보니, 위의 풀이로는 합에 드는 과정을 지워도 10^9 이상의 연산을 하게 될 수 밖에 없다는 걸 깨달았다.

그래서 풀이를 찾아보니 이 문제는 누적합, prefix로 풀이하는 문제였다.

먼저 누적합 배열을 만들어 준다.

위의 예제의 경우 다음과 같은 누적합을 가질 것 이다.

psum = [0, 1, 3, 6, 7, 9]

부분합에서 구간의 합을 구하는 법은, 만약 인덱스 1, 3의 구간의 합을 구하고 싶다면 psum[3] - psum[0]을 하면 된다. 따라서 인덱스 i, j의 구간의 합을 구하고 싶다면, 답은 psum[j] - psum[i - 1]이다.

그렇다면 m으로 떨어지는 구간은 어떻게 구할까?

p[j]와 p[i-1]이 각각 나머지가 같다면, 이 구간은 m으로 나눠 떨어진다.

예를 들어 1 과 7은 둘다 3으로 나눴을때 나머지가 1로 동일하기 때문에 7 - 1 = 6으로 3으로 나누어 떨어진다.

따라서 이 문제는 나머지가 동일한 숫자의 조합의 개수를 구하는 것과 같다고 생각하며 풀이 하면 된다.

정답코드

n, m = map(int, input().split())
numbers = list(map(int, input().split()))
 
remainInfo = [0 for _ in range(m)]
remainInfo[0] = 1
 
total = 0

for number in numbers:
    total += number
    r = total % m
    remainInfo[r] += 1

count = 0

for i in remainInfo:
    count += i * (i - 1) // 2
 
print(count)

 5 3
 1 2 3 1 2


7

코드 해설

remainInfo 라는 배열을 만들어 각 인덱스가 나머지를 가르키게 한 후,

누적합을 total이라는 변수에 담아 구해가며 나머지의 개수를 기록한다.

결국 동일한 나머지의 개수중 두개를 고르는 경우의 수를 구하는 것과 같으므로, i * (i - 1) // 2를 통해 나머지들의 경우의 수를 구해준다.