[Python] 3673번 나눌 수 있는 부분 수열

 

 

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

---

 

24/07/01

 

 

간단한 조합론 + 누적 합 문제이다. 업다운 랜디를 하며 풀었다.

 

문제 풀이 방식 자체는 간단하지만 이를 떠올리는 과정이 좋은 것 같아서 적어보고자 한다.

---

 

문제 접근 방식:

 

 

요구 사항은 연속하는 부분 수열의 합이 $d$로 나누어 떨어지는 것의 개수를 구하는 것이다.(물론 원래 수열 $A$도 주어지고, $d$값도 주어진다.)

 

나의 의식의 흐름은 다음과 같았다.

 

수열의 크기가 테스트 케이스 하나당 최대 5만 정도임.

 

$\rightarrow$ 테스트 케이스 수는 최대 $200$임.

 

$\rightarrow$ 둘을 곱하면 1000만정도 되므로, 하나의 테스트 케이스 당 $\mathcal{O}(N)$의 시간 복잡도로 문제를 해결해야 함. 혹은 하나의 테스트 케이스 당 $\mathcal{O}(N\log N)$의 시간 복잡도로 문제를 해결한다면 간당간당하게 해결할 수 있음.

 

$\rightarrow$ 연속하는 어떤 합을 빠르게 구하는 알고리즘은 누적 합이기 때문에 누적 합을 사용하는 것은 당연해보이고, 누적 합에 DP나 다른 것을 끼얹는 형식으로 갈 것 같음.

 

$\rightarrow$ 따라서 먼저 주어진 테스트 케이스의 누적 합을 구해보자. 어차피 $d$로 나누어 떨어지는 것의 개수만 새면 되기 때문에 누적 합의 모든 원소들에 $\text{mod } d$를 해주어도 상관 없다.

 

$\rightarrow$ $2, 1, 2, 1, 1, 2, 1, 2$의 경우 $0, 2, 3, 1, 2, 3, 1, 2, 0$이 나옴을 확인함.

 

$\rightarrow$ 누적 합 배열에서 연속된 합을 뽑아내기 위해서는 두 수를 뽑고 두 수의 차를 구하면 됨.

 

$\rightarrow$ 생각해보니까 그렇게 구한 구간 합이 $d$로 나누어 떨어지려면 그냥 같은 숫자를 두 개 뽑으면 됨.

$\rightarrow$ 같은 숫자가 $N$개 있을 때, 이 중 $2$개를 뽑는 경우의 수는 $N(N-1)/2$개이므로, 누적 합 배열에서 같은 숫자가 몇 개 있는지 각 숫자마다 세어보자.

 

$\rightarrow$ 근데 카운팅 소트처럼 크기 $d$짜리 배열을 선언하여 세기에는 $d$의 제한이 백만이라 좀 크다.

 

$\rightarrow$ 따라서 딕셔너리로 구현해보자.

 

$\rightarrow$ AC

 

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더보기를 누르면 확인할 수 있다.

# 3673번 나눌 수 있는 부분 수열
# 누적 합, 조합론, 해시맵
'''
요구 사항:
연속하는 부분 수열의 합이 d로 나누어 떨어지는 것의 개수
접근 방법:
수열의 크기는 5만.
테케 수는 200이므로 천만.
아마 하나의 테케당 O(N)의 시복도로 해결해야 할듯?
아니면 O(NlogN)? 간당간당 할듯
나이브하게 생각나는건?
누적합 + 조합론 ㄱㄱ
'''
import sys
input = sys.stdin.readline

def solve():
    D, N = map(int, input().split())
    A = list(map(int, input().split()))
    count = dict()
    total = 0
    count[total] = 1
    for i in range(N):
        total += A[i]
        total %= D
        if total in count:
            count[total] += 1
        else:
            count[total] = 1
    ans = 0
    for key, value in count.items():
        ans += (value*(value-1)//2)
    return ans

C = int(input())
for _ in range(C):
    print(solve())