백준 / 나머지 합 / 10986번 / Python

*문제 출처는 백준에 있습니다.

문제 제목: 나머지 합 / 10986번 (골드 3단계)

문제 사이트: https://www.acmicpc.net/problem/10986

 

문제 설명

---

나의 풀이

def solution(arr):
    answer = 0
    # 구간 합을 넣어둘 배열
    per = [arr[0]]

    if arr[0] % m == 0:
        answer += 1

    for i in range(1, len(arr)):
        if arr[i] % m == 0:
            answer += 1
        per.append(arr[i] + per[i - 1])

    per.pop(0)

    for j in per:
        if j % m == 0:
            answer += 1

    per_len = len(per)

    for k in range(per_len):
        if not per:
            break
        per.pop(0)
        for t in range(len(per)):
            per[t] = per[t] - arr[k]
            if per[t] % m == 0:
                answer += 1

    return answer

n, m = map(int, input().split())

result = list(map(int, input().split()))

print(solution(result))

각 구간별로 하나씩 확인하는 방식을 사용하게되면 시간복잡도는 O(N ^ 2)가 된다. 그래서 시간 초과가 발생한다.

답안

def solution(arr):
    answer = 0
    prefix_sum = 0

    # 나머지를 저장할 리스트 (최대 나머지는 m-1까지 가능하므로 크기는 m)
    mod_count = [0] * m

    # 처음에 나머지가 0인 구간을 카운트하기 위해 0을 미리 설정
    mod_count[0] = 1

    for num in arr:
        prefix_sum += num
        mod = prefix_sum % m

        # 나머지가 음수일 경우를 대비 (파이썬에서 %는 음수 나머지를 반환할 수 있음)
        if mod < 0:
            mod += m

        # 나머지가 같은 경우를 카운트
        answer += mod_count[mod]

        # 나머지를 카운트에 추가
        mod_count[mod] += 1

    return answer


# 입력 받기
n, m = map(int, input().split())
result = list(map(int, input().split()))

# 결과 출력
print(solution(result))

위 문제의 이중 반복문을 하나로 만들어야 한다. 그러므로 나머지를 계산하는 방식을 취했다.

구간 합이 M으로 나누어떨어지려면, 다음 조건을 만족해야 한다.

prefix[j] − prefix[i − 1] ≡ 0 (mod M)

즉, 구간 합이 M으로 나누어떨어지려면, prefix[j]와 prefix[i − 1]의 나머지가 같아야 한다.

prefix[j] ≡ prefix[i−1] (mod M)

---

※ 알아야 할 것

파이썬에서는 나머지가 음수로 등장할 수도 있다.