BOJ 10986 나머지 합

박국현·2022년 5월 5일
0

코테 알고리즘

목록 보기
9/20

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

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

누적합 풀이

문제 구조화

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

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

arr=[A1,,AN]\mathrm{arr} = [A_1, \dots,A_N]

ii번째에서 jj번째 항의 부분구간의 합은 누적합의 차이로 나타낼 수 있다. 0번째 항부터 kk번째 항까지의 누적합을 CkC_k라 하면 부분구간의 합은 다음과 같이 나타낼 수 있다.

Ci1=A1+A2++Ai1Cj=A1+A2++Ai1+Ai++AjC_{i-1} = A_1 + A_2 + \dots + A_{i-1} \\ C_j = A_1 + A_2 + \dots + A_{i-1} + A_i + \dots + A_j \\

위 식의 차이를 구하면

CjCi1=Ai+Ai+1++AjC_j - C_{i-1} = A_i + A_{i+1} + \dots + A_j

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

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

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

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

CjCi1=Mk(kN)C_j - C_{i-1} = M * k \quad (k \in N )

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

CkC_kMM으로 나눈 나머지가 rkr_k라 하자.(0rk<M0 \le r_k < M) 이때 ri1=rjr_{i-1}=r_j를 만족하는 (i,j)(i, j)의 경우의 수는 조합을 이용해 구할 수 있다.
r0,,rnr_0,\dots,r_n을 모두 구한 후 00의 개수를 n0n_0, 11의 개수를 n1n_1,..., kk의 개수를 nkn_k라 하자.

ri1=rj=0 을 만족하는 경우의 수: (n02)ri1=rj=1 을 만족하는 경우의 수: (n12)ri1=rj=M1 을 만족하는 경우의 수: (nM12)모든 경우의 수: (n02)+(n12)++(nM12)r_{i-1}=r_j=0 \text{ 을 만족하는 경우의 수: } {n_0 \choose 2} \\ r_{i-1}=r_j=1 \text{ 을 만족하는 경우의 수: } {n_1 \choose 2} \\ \dots \\ r_{i-1}=r_j=M-1 \text{ 을 만족하는 경우의 수: } {n_{M-1} \choose 2} \\ \text{모든 경우의 수: } {n_0 \choose 2} + {n_1 \choose 2} + \dots + {n_{M-1} \choose 2}

따라서 (n02)+(n12)++(nM12){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 배열의 각 원소가 n0,n1,...,nM1n_0, n_1,...,n_{M-1}이므로 조합을 구하면 답이 된다.

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

전체 코드

논리는 위와 다르지 않지만 코드를 간결하게 해보려고 몇가지를 수정한 최종 코드. 주어진 배열을 따로 변수로 저장하지 않고 풀어봤고, 조합의 합을 구할 때 nk(nk1)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()
profile
공부하자!!

0개의 댓글