백준 10986 | 나머지 합[파이썬 python]

10986번: 나머지 합

 

코드

n, m = map(int, input().split()) # n개의 수, M으로 나누기
A = list(map(int, input().split())) # n개의 수로 이루어진 원본 배열

S = [] # 합 배열
M = [] # 합 배열 S를 M으로 나눈 나머지

answer = 0 # M으로 나눠지는 (i, j) 쌍의 개수

temp = 0
for i in A:
    temp += i
    S.append(temp) # 누적 합
    
    M.append(temp % m) # 누적 합 % m
    
    # 나누어 떨어지면(i에서 i 구간) answer +1
    if temp % m == 0:
        answer += 1
    
# (i, j) 구간 합이 m으로 나누어 떨어지면
# (S[j]-S[i-1]) % m == 0 이므로 S[j] % m == S[i-1] % m 
# 즉, M[j] == M[i-1]이면 (i, j) 구간 합이 m으로 나누어 떨어진다
# 누적합의 나머지가 같은 인덱스 중에 2개를 뽑는 경우의 수 

# 나머지의 개수 파악(m=3이면 나머지는 0, 1, 2 셋 중 하나)
mod_count = [0] * m
for i in M:
    mod_count[i] += 1
    
for i in mod_count:
    answer += i * (i-1) // 2 # i개 중에 2개를 뽑는 경우의 수(2C1)

print(answer)

코드를 이해하고 나면 쉽게 느껴지지만 누적 합에 대한 개념을 확실히 가지고 있어야 풀 수 있는 문제이다.

 

먼저, 합 배열 S와 S를 m으로 나눈 나머지를 저장하는 배열 M을 만든다.

M을 만들 때 누적 합 S가 m으로 나누어 떨어지면(temp % m == 0) (i, i) 구간에서 조건을 만족하므로 순서쌍의 개수 answer+1을 해준다.

S = [] # 합 배열
M = [] # 합 배열 S를 M으로 나눈 나머지

temp = 0
for i in A:
    temp += i
    S.append(temp) # 누적 합
    
    M.append(temp % m) # 누적 합 % m
    
    # 나누어 떨어지면(i에서 i 구간) answer +1
    if temp % m == 0:
        answer += 1

 

이제 나머지 배열인 M에서 m으로 나누어 떨어지는 (i, j) 구간을 구해야 하는데 이를 식으로 보면 다음과 같다. 

(i, j) 구간 합이 m으로 나누어 떨어지면 (S[j]-S[i-1]) % m == 0 이므로 S[j] % m == S[i-1] % m임을 알 수 있다.

즉, S[j] % m = M[j]와 같으므로 식을 다시 쓰면 M[j] == M[i-1]인 구간 (i, j)를 찾아야 하는 것이다.

따라서 합 배열의 나머지가 같은 인덱스 중에 2개를 뽑는 경우의 수를 세어주면 되는데 이를 위해 m으로 나눴을 때의 나머지가 몇 번씩 등장하는지 파악할 배열 mod_count를 만든다.

 

mode_count의 길이는 m과 같다. 그 이유는 m=3이면 3으로 나눴을 때 나머지는 0, 1, 2 셋 중 하나이므로 가능한 나머지는 m개다.

# 나머지의 개수 파악(m=3이면 나머지는 0, 1, 2 셋 중 하나)
mod_count = [0] * m
for i in M:
    mod_count[i] += 1

 

M에서 나머지의 개수를 파악한 뒤에는 경우의 수를 계산해 주면 된다.

2개를 뽑아야 하므로 $_nC_2 = \frac {n \times (n-1)}{2}$임을 활용해 구할 수 있다.

for i in mod_count:
    answer += i * (i-1) // 2 # i개 중에 2개를 뽑는 경우의 수(iC2)

 

answer를 구하기 위해 배열 M과 mod_count 두 개를 사용했는데, 이 과정을 축소해 다음과 같이 더 간결하게 코드를 작성할 수 있다.

n, m = map(int, input().split()) # n개의 수, M으로 나누기
A = list(map(int, input().split())) # n개의 수로 이루어진 원본 배열

S = [] # 합 배열
mod_count = [0] * m

temp = 0
for i in A:
    temp += i
    S.append(temp) # 누적 합
    mod_count[temp % m] += 1 # 누적 합 % m
    
# m으로 나누어 떨어지는 (i, j) 순서쌍의 개수
answer = mod_count[0] # 나누어 떨어지는 경우 = (i, i)
for i in mod_count:
    answer += i * (i-1) // 2 # i개 중에 2개를 뽑는 경우의 수(iC2)

print(answer)