BOJ 10986 나머지 합

처음에는 길이가 1인 부분 구간, 2인 부분 구간...해서 길이가 $n$인 부분 구간인 경우를 모두 구하려 했다. 당연히 시간 초과가 났다.

찾아보니 누적합과 나머지를 사용하여 풀 수 있는 문제였다. 예시 입력인 1 2 3 1 2을 넣었을 때의 풀이과정을 보자.

누적합 풀이

문제 구조화

누적합으로 부분구간의 합 나타내기

문제에서 주어진 수열을 $A_i(i=1,...,N)$이라 하자.

$i$번째에서 $j$번째 항의 부분구간의 합은 누적합의 차이로 나타낼 수 있다. 0번째 항부터 $k$번째 항까지의 누적합을 $C_k$라 하면 부분구간의 합은 다음과 같이 나타낼 수 있다.

위 식의 차이를 구하면

즉 $C_j - C_{i-1}$를 통해 $i$번째 항부터 $j$번째 항까지 부분 구간의 합을 구할 수 있는 것이다.

참고: $C_0 = 0$, 1번째 항을 포함하는 부분구간의 합을 구하기 위해 필요함

누적합의 나머지로 답 구하기

이 문제는 부분구간의 합이 $M$으로 나누어 떨어지는 경우를 구해야 한다. 즉 $C_j - C_{i-1}$를 $M$으로 나누었을 때 나머지가 $0$이어야 하는 것이다.

이를 만족하기 위해서는 $C_j$와 $C_{i-1}$ 두 수가 $M$으로 나눈 나머지가 일치해야 한다. 다시 말해 $M$으로 나눈 나머지가 일치하는 $C_{i-1}$와 $C_j$를 만족하는 $(i, j)$의 개수가 이 문제의 답이 되는 것이다.

$C_k$를 $M$으로 나눈 나머지가 $r_k$라 하자.($0 \le r_k < M$) 이때 $r_{i-1}=r_j$를 만족하는 $(i, j)$의 경우의 수는 조합을 이용해 구할 수 있다.
$r_0,\dots,r_n$을 모두 구한 후 $0$의 개수를 $n_0$, $1$의 개수를 $n_1$,..., $k$의 개수를 $n_k$라 하자.

따라서 ${n_0 \choose 2} + {n_1 \choose 2} + \dots + {n_{M-1} \choose 2}$을 구하는 것이 이 문제의 답이 되는 것이다.

코드로 풀기

문제에서 주어진 수열을 누적합 수열로 바꿔주자.

arr = list(map(int, input().split())
for i in range(1, len(arr)):
	arr[i] += arr[i-1]
# arr: [1, 3, 6, 7, 9]

배열의 제일 앞에서 새로운 원소 0을 추가하고 각 원소를 M(문제에서 준 나누는 수)으로 나눈 나머지로 바꿔주자.

arr = [0] + arr
for i in range(len(arr)):
	arr[i] %= M
# arr: [0, 1, 0, 0, 1, 0]

0부터 M-1까지 각 숫자가 몇 개 들어있는지 셀 cnt 배열을 만들고 기록을 하자.

cnt = [0] * M
for i in range(len(arr)):
	cnt[arr[i]] += 1
# cnt: [4, 2, 0]

cnt 배열의 각 원소가 $n_0, n_1,...,n_{M-1}$이므로 조합을 구하면 답이 된다.

answer = 0
for num in cnt:
	answer += num * (num-1) // 2

전체 코드

논리는 위와 다르지 않지만 코드를 간결하게 해보려고 몇가지를 수정한 최종 코드. 주어진 배열을 따로 변수로 저장하지 않고 풀어봤고, 조합의 합을 구할 때 $n_k * (n_k-1)$의 합을 먼저 구하고 마지막에 2로 나눠줬다.

import sys

input = sys.stdin.readline
print = lambda x: sys.stdout.write(f'{x}\n')


def main():
    N, M = map(int, input().split())
    remainders = [0] * M
    remainders[0] = 1
    num = 0
    for new_num in map(int, input().split()):
        num = (new_num + num) % M
        remainders[num] += 1

    answer = sum(i * (i - 1) for i in remainders) // 2
    print(answer)


if __name__ == '__main__':
    main()