[백준 | BOJ] 10986번: 나머지 합 (Python)

문제

https://www.acmicpc.net/problem/10986

 

난이도

 

알고리즘 분류

• 수학
• 누적합

 

요약

• 수학적 지식이 필요한 문제
• 난이도에 비해 코드는 쉽게 짤 수 있음(=풀이 과정 자체는 별거 없음)
• 하지만 아이디어를 떠올리는 것이 어려웠던 문제

---

문제 접근 방식

어떤 알고리즘을 사용해야 하는지 살펴보기 위해 문제의 조건들을 살펴보자. 3번에 대한 내용은 쉽게 떠올리기 힘들어서 체감 난이도가 높은 것 같다.

1. 연속된 부분 구간의 합을 구해야 한다.
   
   • 누적 합 알고리즘을 사용해서 구할 수 있다.
   • 시간 복잡도는 대략 O(N)
2. N의 범위가 최대 1,000,000 이다.
   • 허용 가능한 최대 시간 복잡도는 대략 O(N log N) 임을 알 수 있음
   • 누적 합 알고리즘은 O(N) 이므로 사용 가능!
3. 구간 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 모두 구해야 한다.
   • 우선 모듈러(modular) 연산의 성질을 알아야 한다.
   • 두 수 A, B가 있을 때 (A - B)%M = 0 이라면 (A%M) - (B%M) = 0 이다. 따라서 A%M = B%M 이다.

---

풀이 과정

1. 우선 누적 합 알고리즘을 사용해 연속된 부분 구간의 합을 표현하는 배열을 만든다.
   • 주어진 배열을 A라고 하고, 누적 합 배열을 psum이라고 하면 아래와 같이 쉽게 구할 수 있습니다.
     

     

   • 저는 누적 합 배열을 생성할 때 psum[i] = 앞에서부터 i개의 원소의 합으로 생각하면 이해하기 더 쉬운 것 같아서 psum[0] = 0으로 설정하며, N+1 길이의 배열로 정의하는 것을 선호하는 편입니다.
     
2. 이제 구간 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 모두 구해야 한다.
   • 여기서 (i, j)는 배열 A의 인덱스를 기준으로 한다.
   • 문제 접근 방식을 통해 모듈러 연산의 성질을 알아보았고, 이를 문제에 적용해보자.
   • (i <= j)를 만족하는 임의의 i, j를 선택 했을 때, 연속된 부분 구간의 합은 psum[j+1] - psum[i] 이다. 
     • psum[j+1] = A[0] + A[1] + ... + A[j] (앞에서부터 j+1개의 원소)
     • psum[i] = A[0] + A[1] + ... + A[i-1] (앞에서부터 i개의 원소)
     • psum[j+1] - psum[i] = A[i] + A[i+1] + ... + A[j] (i ~ j 까지의 연속된 부분 구간의 합)
   • 구해야 할 답은 (psum[j+1] - psum[i])%M = 0 을 만족하는 모든 (i, j) 쌍의 개수이다.
   • 식을 정리하면, psum[j+1]%M = psum[i]%M 을 만족하는 (i, j) 쌍의 개수를 구하는 것과 같다.
   • 즉, M으로 나누었을 때 나머지가 같은 (i, j) 쌍의 개수를 구하는 것과 같다.
     
     
3. 나머지가 같은 (k, l) 쌍을 구해야 하므로 이를 계산하는 배열을 추가로 생성한다.
   • 여기서 (k, l)은 배열 psum의 인덱스를 기준으로 한다.
   • 구간 합 배열에서 M으로 나누었을 때의 나머지를 index로 갖는 배열을 생성한다. (최대 M-1 길이의 배열)
   • psum[0]은 원소가 없는 경우이므로 제외하고 psum[1] ~ psum[N]에 대해서만 구한다.
   • 예제(M=3)를 기준으로 아래와 같은 배열이 생성된다.

     

   • 나머지(index) 0: psum[2], psum[3], psum[5]
     • 나머지가 0인 경우는 각각 나누어 떨어지므로 따로 카운팅 해줘야 함. 자세한 내용은 4번에서 설명
   • 나머지(index) 1: psum[1], psum[4]
   • 나머지(index) 2: 없음
     
     
4. 나머지가 같은 (k, l) 쌍의 개수를 구한다.
   • nCr = n * (n-1) / 2 임을 활용한다.
   • 나머지가 0인 개수 : 3
     • 3개의 원소중 서로 다른 (k, l) 쌍을 뽑는 경우의 수 = 3C2 = 3 * 2 / 2 = 3
   • 나머지가 1인 개수 : 2
     • 2개의 원소중 서로 다른 (k, l) 쌍을 뽑는 경우의 수  = 2C2 = 2 * 1 / 2 = 1
   • 여기서 주의할 점은 나머지가 0인 경우는 그 자체로도 나누어 떨어지므로 따로 카운팅을 해야 한다.
     • k = 0 인 경우에 해당하며, 이 경우는 remainder[0]의 개수와 같다.
     • 예제를 기준으로 각각 psum[2], psum[3], psum[5]에 해당함.
   • 위 세가지 경우를 모두 더하면 3 + 1 + 3 = 7
     
     
5. 정답 검증
   • A = [1, 2, 3, 1, 2]
   • 나머지가 0인 개수를 구하는 과정을 살펴보자.
     • 서로 다른 (k, l) 쌍 : [2, 3], [2, 5], [3, 5]
       • k = 2, l = 3 인 경우 psum[3] - psum[2] = 3 이므로 A[2]인 경우 (3)
       • k = 2, l = 5 인 경우 psum[5] - psum[2] = 6 이므로 A[2]+A[3]+A[4]인 경우 (3+1+2)
       • k = 3, l = 5 인 경우 psum[5] - psum[3] = 3 이므로 A[3]+A[4]인 경우 (1+2)
     • k = 0인 경우: [0, 2], [0, 3], [0, 5]
       • k = 0, l = 2 인 경우 psum[2] - psum[0] = 3 이므로 A[0]+A[1]인 경우 (1+2)
       • k = 0, l = 3 인 경우 psum[3] - psum[0] = 6 이므로 A[0]+A[1]+A[2]인 경우 (1+2+3)
       • k = 0, l = 5 인 경우 psum[5] - psum[0] = 9 이므로 A[0]+A[1]+A[2]+A[3]+A[4]인 경우 (1+2+3+1+2)
   • 나머지가 1인 개수를 구하는 과정을 살펴보자.
     • 서로 다른 (k, l) 쌍: [1, 4]
     • k = 1, l = 4인 경우 psum[4] - psum[1] = 6 이므로 A[1]+A[2]+A[3]인 경우 (2+3+1)

---

코드 설명

# 구간 합 배열 구하기
psum = [0] * (N + 1)
for i in range(N):
    psum[i + 1] = psum[i] + arr[i]

구간 합 배열을 구하는 부분입니다. psum[0] = 0(원소를 선택하지 않은 경우)이고, 구간 합 배열은 psum[1] 부터 계산합니다.

사실 psum 배열을 따로 생성하지 않고 아래 부분과 같이 값만 저장하여 사용할 수도 있습니다.

 

# 나머지 배열 구하기
remainder = [0] * M
for i in range(1, N + 1):
    remainder[psum[i] % M] += 1

################################

# psum 배열 사용x
remainder = [0] * M
psum = 0
for i in range(N):
    psum += arr[i]
    remainder[psum % M] += 1

구간 합 배열을 M으로 나누었을 때 나머지의 개수를 구하는 부분입니다.

원소를 선택하지 않은 psum[0] 부분을 제외하고 psum[1]부터 나머지를 계산합니다.

 

# M으로 나누어 떨어지는 개수 세기
answer = remainder[0]  # 나머지가 0인 개수
for i in range(M):
    answer += (remainder[i] * (remainder[i] - 1)) // 2

나머지가 0인 경우는 이미 나누어 떨어지므로 개수를 더해주고, nCr = n * (n-1) / 2를 활용하여 계산해 줍니다.

---

전체 코드

import sys

input = sys.stdin.readline


def main() -> None:
    N, M = map(int, input().rstrip('\n').split(' '))
    arr = list(map(int, input().rstrip('\n').split(' ')))

    # 구간 합 배열 구하기
    psum = [0] * (N + 1)
    for i in range(N):
        psum[i + 1] = psum[i] + arr[i]

    # 나머지 배열 구하기
    remainder = [0] * M
    for i in range(1, N + 1):
        remainder[psum[i] % M] += 1

    # M으로 나누어 떨어지는 개수 세기
    answer = remainder[0]  # 나머지가 0인 개수
    for i in range(M):
        answer += (remainder[i] * (remainder[i] - 1)) // 2

    print(answer)


if __name__ == "__main__":
    main()