[골드 3] 백준 3673 - 나눌 수 있는 부분 수열 (파이썬)

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

풀이

• 길이 N인 수열이 있을 떄,
• 부분 수열의 합이 D로 나누어 떨어지는 경우의 수를 구해라

누접합 배열을 만들어서, 누적합을 D로 나눈 나머지를 활용하면 풀이가 가능하다.

T = int(input())

for _ in range(T):
    D, N = map(int, input().split())
    A = list(map(int, input().split()))

    psum = [0] * N
    psum[0] = A[0]
    for i in range(1, N):
        psum[i] = psum[i-1] + A[i]

누적 합 계산을 진행한다.

테스트 케이스 T개 만큼 반복하면서, 각 테스트 케이스마다 나누는 수 D와 배열 크기 N을 입력 받고, 배열 A를 입력받는다.

그 다음 psum[i] = 0 ~ i 까지 합을 계산해서 저장한다.

    answer = 0
    count = {}
    count[0] = 1

답을 담을 answer를 0으로 초기화하고, 나머지를 키로, 나머지가 몇번 나왔는지 저장하는 count 딕셔너리를 선언한다.

왜 카운트 0은 1로 초기화를 하는가? → 처음부터 psum[i]자체가 D로 나누어 떨어지는 경우를 바로 세기 위해서이다.

    for i in range(N):
        if psum[i] % D not in count:
            count[psum[i] % D] = 0
        answer += count[psum[i] % D]
        count[psum[i] % D] += 1
    
    print(answer)

배열을 순회하면서, 각 psum[i]에 대해, psum[i] % D를 구하고, 해당 나머지가 count에 없으면 0으로 초기화 한다.

지금까지 같은 나머지를 가진 누적합이 몇번 나왔는지를 answer에 더하고 출력한다.

전체 요약의 흐름은 다으모가 같다.

1. 입력 받고
2. 누적합 배열 만든 다음
3. 누적합을 D로 나눈 나머지를 기준으로 개수 세서
4. 같은 나머지가 있을 때마다 답을 올린다

코드

# 백준 3673 - 나눌 수 있는 부분 수열
# 분류 : 누적 합

T = int(input())

for _ in range(T) :
    D, N = map(int, input().split())
    A = list(map(int, input().split()))

    psum = [0] * N
    psum[0] = A[0]
    for i in range(1, N) :
        psum[i] = psum[i - 1] + A[i]
    
    answer = 0
    count = {}
    count[0] = 1 # psum[i] % D == 0 0~i
    for i in range(N) :
        if psum[i] % D not in count :
            count[psum[i] % D] = 0
        answer += count[psum[i] % D]
        count[psum[i] % D] += 1
    
    print(answer)