BOJ 21268 Do Use FFT

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

솔루션에서 Tellegen's Principle을 인용하고 있어서 지난 포스트와 연관해 작성해본다. 머리 좋은 사람들은 아마 생각하지 않고도 풀 수 있을 것 같다.

Without Tellegen's principle

$E_{k} = \sum_{i = 1}^{n} C_{i} \prod_{j = 1}^{k} (A_{i} + B_{j})$를 답이라고 하자.

$i$와 관련한 항들을 사부작거리다보면, $P_{j} = \sum_{i = 1}^{n} C_{i}A_{i}^{j}$를 미리 계산해두면 편할 것 같다.

$\sum_{j} F_{k, j}x^{j} = (x+B_1)(x+B_2)\cdots (x+B_{k})$라고 두면, $E_{k} = \sum_{j} P_{j} F_{k, j}$가 된다는 사실을 알 수 있다. 만약 $F_{k, k-j}$꼴이었다면 전형적인 다항식 곱하기인데 미묘하게 다르다. 그럴때는 $x \leftarrow x^{-1}$을 대입해주고 $x^{k}$를 곱해서 $G_{k}(x) = (B_1 x + 1) \cdots (B_k x + 1)$을 생각해주자.

그러면 $P(x) = \sum_{j} P_{j}x^{j}$를 생각했을 때 $E_{k} = [x^k] P(x)G_{k}(x)$가 된다. 일단 $G_{1}(x), \cdots, G_{n}(x)$를 다 곱해볼 수는 당연히 없고, 각 $k$마다 $x^k$의 계수만 필요하다는 것에 주목해 분할 정복을 생각해보자.

구간 $[s, e]$에 대해 $E_{s}, \cdots, E_{e}$를 구해야 한다고 생각해보자. $P(x)G_{s-1}(x)$를 미리 알고 있다고 가정하고, $[s, m]$에 대한 oracle을 호출해 $E_{s}, \cdots, E_{m}$까지는 알고 있다고 가정한다. 이제 $P(x)G_{s-1}(x)$에 $(B_s x + 1) \cdots (B_m x + 1)$을 곱해서 $P(x)G_{m}(x)$를 만들고 $[m+1, e]$로 분할정복을 하고 싶다.

문제는 $P(x)$의 차수가 $n$까지 커질 수 있으니 분할 정복이 안 된다. 연산량이 적어도 $e-s$에 대한 함수로 나와야 할 텐데..

그런데 잘 생각해보면 $i \in [m+1, e]$에 대해 $P(x)G_{i}(x)$의 $i$차항은 $(B_{m+1} x + 1) \cdots (B_{i} x + 1)$이 $i - m$차식임을 고려할 때 $P(x)G_{m}(x)$의 $m$차항 이상 $e$차항 이하만 확인하면 된다는 사실을 알 수 있다.

그렇다면 $P(x)G_{s-1}(x)$를 다 들고 다니는 대신, $P(x) G_{s-1}(x)$의 $[s, e]$차항만 들고 다니자. 이는 $e - s+1$차식으로 충분히 감당 가능하다. $[s, m]$ 구간을 분할정복하고, $m-s$차식인 $(B_s x + 1) \cdots (B_m x + 1)$을 곱하고, $[m+1, e]$차항만 뽑아내서 $[m+1, e]$ 구간을 분할정복해주면 된다. $\mathcal{O}((e - s)\log (e - s))$ 정도의 연산이 필요하니 전체 시간 복잡도는 $\mathcal{O}(N \log^{2} N)$이다.

$P(x)$를 어떻게 계산하는지 설명 안 했는데, 결국 $P(x) = \sum_{i = 1}^{n} \frac{C_i}{1 - A_{i} x}$꼴이다. $C_{i} = 1$인 경우에는 적분식이 $\log \prod_{i = 1}^{n} (1 - A_{i}x)$ 꼴인 걸 이용해서 곱하고 로그취하고 미분하는 방식이 있는데, 사실 그냥 더 간단하게 분할 정복으로 분자, 분모를 유지하면서 더해주면 역시 $\mathcal{O}(N\log^{2} N)$이면 된다.

With Tellegen's Principle

$E_{k} = \sum_{j} P_{j} F_{k, j}$ 에서, $F_{k, j}$를 행렬로 보는 것이 자연스럽다. 그렇다면 이 transpose는 어떨까?
벡터 $a_{1}, \cdots, a_{N}$이 주어졌을 때 $b_{j} = \sum_{k} a_{k} F_{k, j}$를 구한다고 생각해보면, $\sum_{j} b_{j} x^{j} = \sum_{k} a_{k}(x + B_1) \cdots (x + B_k)$가 되는 걸 알 수 있다. 저 다항식 곱 계산은 분할 정복으로 쉽게 할 수 있다.

Tellegen's principle을 적용하여 해당 알고리즘의 transpose를 구해보자. 가장 자연스러운 분할 정복 알고리즘은 다음과 같다.

def solve([s, e], a) -> polynomial[x]:
    if s == e:
        return a[s] # constant polynomial

    m = (s + e) / 2
    p = solve([s, m], a)
    B = (x + B_s) ... (x + B_m) # acquired from another DnC
    q = solve([m+1, e], a)

    return p + B * q

즉 solve(s, e, a)라는 연산자는 $e-s+1$개의 수 $a_s, \cdots, a_e$를 입력으로 받아 똑같이 $\sum_{j} \beta_{j}x^{j} = p + Bq$를 만족하는 $\beta_s, \cdots, \beta_e$로 돌려주는 linear operator로 생각할 수 있다. 딱딱하게 쓰면,

$$\mathrm{solve}[s, e] = \mathrm{solve}[s, m] + \mathrm{multiply}_{B} \circ \mathrm{solve}[m+1, e]$$

이 transpose는 다음과 같다.
$$\mathrm{solve}^{\mathbf{t}}[s, e] = \mathrm{solve}^{\mathbf{t}}[s, m] + \mathrm{solve}^{\mathbf{t}}[m+1, e] \circ \mathrm{multiply}_{B}^{\mathbf{t}}$$

지난 포스트에서 multiply의 transpose가 middle product라고 썼는데, 나한테는 이것저것 다 헷갈려서 그냥 $B(x^{-1})$을 곱하고 non-negative 차수만 남긴 거라고 외운다. 이 연산이 어떻게 생겼는지 구현하다보면 위의 풀이와 완전히 동일한 결론을 얻는다.

여담으로 $P(x)$를 계산하는 과정 역시 dual이 있는데, $C_{i}$가 주어졌을 때 $P(x)$를 계산하는 것을 generalized power sum이라고 하고, 그 transpose는 multiple point evaluation이 된다.