[Baekjoon]백준 10986 나머지 합(골드 3) - Python

문제설명

문제를 살펴보면 부분합 중에서 입력받은 수로 나누었을 때 나머지가 0인 경우의 수를 출력해야한다

일단 혹시나 싶은 마음에 누적 합, 부분 합을 따로 구하여 나머지가 0이 되는 경우의 수를 찾는 코드를 작성해보았다

num, div = map(int, input().split())
ls = list(map(int, input().split()))
cnt = 0

prefix_sum = [0] * (num + 1) # 누적 합 계산
for i in range(num):
    prefix_sum[i+1] = prefix_sum[i] + ls[i]

result = []
for i in range(num): # 구간 합 계산
    for j in range(i, num):
        result.append(prefix_sum[j+1] - prefix_sum[i])

for i in result:
    if i % div == 0:
        cnt += 1

print(cnt)

해당 코드를 제출하였더니 메모리 초과가 발생하였다

음...? 메모리 초과? 그래서 들었던 생각이 result 리스트의 크기를 대폭 줄여보자! 였다

result리스트의 크기를 줄이는 방법으로는 부분합을 계산함과 동시에 나머지 연산을 하여 

그 값이 0이라면 cnt 변수의 값을 증가시키는 방법으로 해보았다

안타깝게도 입력받는 수가 많다보니 시간 초과가 발생하는 것을 확인할 수 있었다

num, div = map(int, input().split())
ls = list(map(int, input().split()))
prefix_sum = [0] * (num + 1)
cnt = 0

for i in range(num): # 누적 합 계산
    prefix_sum[i+1] = prefix_sum[i] + ls[i]

for i in range(num): # 구간 합 계산
    for j in range(i, num):
        if (prefix_sum[j+1] - prefix_sum[i]) % div == 0:
            cnt += 1

print(cnt)

그래서 이번에는 누적합의 각 나머지 값을 저장한 후에 이를 활용하는 방법으로 시도해보았다

이때까지는 for문을 중첩하여 사용하였기 때문에 시간초과가 발생했을거리라 생각하고

하나의 for문으로 해결해보기 위해서 시도해본 방법이었다

 

일단, 누적합을 구하여 리스트에 저장을 함과 동시에 나머지를 index로 가지는 리스트에 +1을 해준다

(나머지를 index로 가진다 → 나머지가 3인 경우 remainder[0], remainder[1], reamainder[2])

부분합을 구하는 방법은 이전 문제에서도 해보았듯이 각 index의 누적합을 서로 빼주면 된다

 

여기서부터가 중요한데, 예제의 경우로 생각해보자면 remainder에는 1 0 0 1 0이 저장되어있을것이다

그렇다면 나머지가 0인 경우는 3가지, 1인 경우가 2가지일텐데, 나머지가 1인 경우를 서로 빼면 나머지가 0이 된다

같은 원리로 나머지가 2인 경우도 서로 값을 빼면 나머지가 0이 된다

예를 들어, 2와 8이 있을 때 각 수는 3으로 나누었을 때 나머지가 2이지만, 8-2를 한 후에는 나머지가 0이 된다

이를 활용하여 조합으로 각 경우의 수를 구해줄 수 있다

각 나머지의 개수에서 두 개의 index를 뽑는 것은 nC2와 동일하기 때문이다

 

예제의 경우로 다시 돌아와서 생각해보자면 나머지가 0인 경우는 굳이 두 개의 index를 고르지 않아도

나머지가 0이기 때문에 초기 값을 remainder[0]으로 설정해준다 즉, 3을 초기값으로 설정한다

이후에 나머지가 0인 경우의 index가 3이므로 3C2를 해주면 3이 되고, 나머지가 1인 경우의 index는 2이므로

2C2를 하면 1, 나머지가 2인 경우의 index는 0이기 때문에 이 값들을 다 더해주면 3 + 3 + 1이 되어 7이 된다

이 방법으로 코드를 제출하면 정답임을 확인할 수 있다

num, div = map(int, input().split())
ls = list(map(int, input().split()))
remainder = [0] * div # 나머지값 저장

prefix_sum = 0
for i in range(num):
    prefix_sum += ls[i]
    remainder[prefix_sum % div] += 1

result = remainder[0] # 나머지가 0이 되는 경우의 수
for i in remainder:
    result += i * (i-1) // 2
    
print(result)