[백준] 10986번: 나머지 합 (Python) - 메모리 초과와 누적합

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

10986번: 나머지 합

1. 첫번째 시도 -> 메모리 초과 발생

256MB 

• int형 == 4Byte
• 1KB == 1024Byte
• 1MB == 1024KB
• 256MB = 256 * 1024KB = 256 * 1024 * 1024B = int형 256 * 1024 * 1024 / 4Byte = 67108864개

 

[아이디어]

• 2차원 배열 dp 를 생성하고 여기에 i부터 j 까지의 누적합을 저장한다. 
• 저장된 누적합 2차원 배열을 하나씩 순회하면서 M으로 나눈 나머지를 확인한다. 

# https://www.acmicpc.net/problem/10986
N, M = map(int, input().split())
data = list(map(int, input().split()))
answer = 0
dp = [[0] * (N+1) for _ in range(N+1)] # 누적합 저장 dp[i][j] 는 i ~ j 까지의 합

# 누적합 저장하기
for i in range(N):
    for j in range(i, N):
        if i == j:
            dp[i][j] = data[i]
        else:
            dp[i][j] = dp[i][j-1] + data[j]

# check dp
# for i in range(N):
#     print(*dp[i])

# 누적합을 M 으로 나누기
for i in range(N):
    for j in range(N):

        if dp[i][j] != 0 and dp[i][j] % 3 == 0:
            answer += 1
print(answer)

 

2. 두번째 시도 -> 시간 초과

[아이디어]

• dp 를 1차원배열으로 바꾸고, 누적합을 사용하여 i ~ j 구간의 합을 구하자 !
• 메모리 초과는 방지했으나 시간초과가 발생했음. 현재 구간 누적합을 O(n**2) 시간복잡도로 구하고 있어서 시간 초과가 발생한다. 

# https://www.acmicpc.net/problem/10986
import sys

input = sys.stdin.readline
N, M = map(int, input().split())
data = list(map(int, input().split()))
answer = 0
dp = [0] * N
dp[0] = data[0]
for i in range(1, N):
    dp[i] = dp[i-1] + data[i]

# 누적합을 M 으로 나누기,
for i in range(N-1):
    for j in range(i, N):
        prefix_sum = dp[j] - dp[i-1]
        if prefix_sum % 3 == 0:
            answer += 1

 

[최종 풀이]

풀이를 참고해보니 몫들을 누적합으로 저장해서 직접 나누는 것이 아니라 '나머지'들만 따로 저장하여 나머지가 0 인 구간을 구하는 방식으로 푸는 아주아주 참신한 풀이가 있었다 ... 

# https://www.acmicpc.net/problem/10986
import sys

input = sys.stdin.readline
N, M = map(int, input().split())
data = list(map(int, input().split()))
answer = 0
dp = [0] * M # 나머지를 저장하는 배열 , size가 M 이 된다.
prefix_sum = 0
for i in range(N):
    prefix_sum += data[i]
    dp[prefix_sum%M] += 1 # 나머지가 몇개인지 저장하기

# 누적합을 M 으로 나누기,
answer = dp[0] # 나머지가 0 인 경우는 무조건 포함시키기 

# iC2 조합으로 구하기 , 나머지가 같은 구간 2개를 뽑는 모든 경우의 수 
for i in dp:
    answer += i*(i-1)//2 
print(answer)

 

참고 자료 

https://book.acmicpc.net/algorithm/prefix-sum

누적 합