AtCoder

[ABC221 H] Count Multiset

IBory 2023. 11. 16. 18:01

https://atcoder.jp/contests/abc221/tasks/abc221_h
kenkoooo 레이팅: 2793
체감 난이도: D5
Tag: DP, Combinatorics


두 수 $N, M$이 주어진다. $1$ 이상 $N$ 이하의 모든 정수 $k$에 대해, 아래 조건을 만족하는 집합의 수를 $998\ 244\ 353$로 나눈 나머지를 구해야 한다.

  • 집합에 있는 원소의 수는 $k$이다.
  • 모든 원소의 합은 $N$이다.
  • 같은 원소는 최대 $M$개까지 포함할 수 있다.

집합은 순서를 구분하지 않기 때문에 막 세는 것이 조금 까다롭다. 핵심은 구하고자 하는 것을 잘 변형해서 세기 쉬운 형태로 잘 바꿔주는 것이다. 풀이에서는 순서를 구분해서 세도 되는 형태로 조건을 잘 변형한다.
 
$A$는 항상 오름차순으로 정렬되어 있다고 가정한다. 그리고 인접한 $A$의 차이가 합에 기여하는 정도를 대신 센다고 하자. 이러면 값의 순서에 따라 합에 기여하는 정도가 달라지므로, 순서라는 조건을 값으로 치환해서 생각할 수 있다.
$A_0 = 0$이라고 생각하고, $B_i = A_i - A_{i - 1}$라고 하자. 그러면 구하고자 하는 건 아래 조건을 만족하는 $B$의 배열이다.

  • $B$에서 연속된 $0$의 개수는 최대 $M-1$개 이다. (같은 수가 $M$개 이상 등장할 수 없으므로)
  • $B_1$은 양수이다. ($A_i$가 양수여야 하므로)
  • $\sum_{i=1}^{k} {B_i \times \left( k - 1 + i \right) } = N$

 
모든 $k$에 대해 답을 구해야 하는 입장에서 세 번째 조건이 처리하기 까다롭다. 대신 여기서 인덱스를 뒤집어서 수를 센다고 하자. 그러면 두 번째 조건은 $B_k$가 양수여야 한다는 조건으로, 세 번째 조건의 식은 $ \sum_{i=1}^{k} {B_i \times i } = N$ 와 같이 변하게 된다.
 
이제 아래와 같이 dp 상태를 정의해서 문제를 해결할 수 있다.
$dp[n][s]$: $B$의 $n$번째 값까지 채웠고, 채운 원소의 합이 $s$가 되는 경우의 수
세 번째 조건에 따라 $n$번째로 원소를 넣을 때는 $n, 2n, \cdots$ 에 해당하는 원소만 넣을 수 있다. 또, 조건을 편하게 다루기 위해 항상 양수인 값만 넣는다고 하자. $B$에 $0$이 들어가는 경우는 따로 처리해 줄 것이다. 이를 이용해 점화식을 세워보면 아래와 같다.
$$dp[n][s] = \sum_{k = 1}^{ \lfloor \frac{s}{n} \rfloor } \sum_{t = 1}^{t \le \min(t, M -1)} dp[n-t][s-kn]$$
시그마가 두 개 있어서 대충 생각하면 $O(N^3M)$ 같지만, 잘 생각해 보면 누적합으로 잘 처리할 수 있는 형태이다. 따라서 각 상태 계산을 $O(1)$에 할 수 있고, 전체 문제를 $O(N^2)$에 해결할 수 있다.
 
https://atcoder.jp/contests/abc221/submissions/47596249

'AtCoder' 카테고리의 다른 글

[ABC254 Ex] Multiply or Divide by 2  (5) 2023.11.19
[ABC241 Ex] Card Deck Score  (6) 2023.11.14
[ABC259 Ex] Yet Another Path Counting  (5) 2023.11.12
[ABC267 Ex] Odd Sum  (1) 2023.11.12