[codeforces] 1921F. Sum of Progression

starbow·2024년 6월 7일

PS/CP

목록 보기

10/25

길이가 $n$ 인 정수 배열 $a$ 와 $q$ 개의 쿼리가 $s, d, k$ 형식으로 주어집니다. 각 쿼리에 대해서 $a_{s} + 2 \cdot a_{s+d} + \cdots + k \cdot a_{s + d \cdot (k-1)}$ 을 구하면 됩니다.

$n$ 은 최대 10만, $q$ 는 최대 20만까지 주어지기 때문에 각 쿼리에 대해서 Naive하게 계산하면 TLE를 받습니다. (시간복잡도 $O(nq)$ ) 시간복잡도를 줄일 수 있는 다른 방법을 생각해봐야 합니다.

잘 생각해보면 $d$ 가 클 수록 더하는 값의 갯수가 줄어들 수 밖에 없다는것을 알 수 있습니다. 즉, $d$ 가 클 때는 그냥 Naive하게 더하고 $d$ 가 작을때는 다른 방식으로 구하는 방향으로 풀이를 생각해 볼 수 있습니다.

그럼 케이스를 나눠서 생각해보겠습니다.

Case 1) $d \geq \sqrt{n}$

이 때는 그냥 Naive하게 더합니다. 항의 갯수가 $O(\sqrt{n})$ 이므로 $O(\sqrt{n})$ 에 값을 구할 수 있습니다.

Case 2) $d < \sqrt{n}$

$a_{0} = 0$ 으로 설정하면 식을 아래와 같이 변형시킬 수 있습니다.

\begin{aligned} &a_{s} + 2 \cdot a_{s+d} + \cdots + k \cdot a_{s + d \cdot (k-1)} \\ = \, &a_{(s \, \mathrm{mod} \, d)} + 2 \cdot a_{(s \, \mathrm{mod} \, d)+d} + \cdots + \left\lfloor{\frac{s}{d}}\right\rfloor \cdot a_{s-d} + \left(\left\lfloor{\frac{s}{d}}\right\rfloor+1 \right) \cdot a_{s} + \cdots + \left(\left\lfloor{\frac{s}{d}}\right\rfloor+k \right) \cdot a_{s + d \cdot (k-1)} \\ &- \left( a_{(s \, \mathrm{mod} \, d)} + 2 \cdot a_{(s \, \mathrm{mod} \, d)+d} + \cdots + \left\lfloor{\frac{s}{d}}\right\rfloor \cdot a_{s-d} \right) \\ &- \left\lfloor{\frac{s}{d}}\right\rfloor \cdot \left( a_{s} + a_{s+d} + \cdots + a_{s + d \cdot (k-1)} \right) \\ \end{aligned}

그리고 $A_{s, d}$ 와 $B_{s, d}$ 를 아래와 같이 정의 하겠습니다.

\begin{aligned} &A_{s, d} = a_{(s \, \mathrm{mod} \, d)} + 2 \cdot a_{(s \, \mathrm{mod} \, d)+d} + \cdots + \left\lfloor{\frac{s}{d}}\right\rfloor \cdot a_{s-d} + \left(\left\lfloor{\frac{s}{d}}\right\rfloor+1 \right) \cdot a_{s} \\ &B_{s, d} = a_{(s \, \mathrm{mod} \, d)} + a_{(s \, \mathrm{mod} \, d)+d} + \cdots + a_{s-d} + a_{s} \end{aligned}

그럼 우리가 구해야 하는 식은 아래와 같이 변형됩니다.

\begin{aligned} &a_{s} + 2 \cdot a_{s+d} + \cdots + k \cdot a_{s + d \cdot (k-1)} \\ = \, &A_{s + d \cdot (k-1), d} - A_{s-d, d} - \left\lfloor{\frac{s}{d}}\right\rfloor \cdot \left( B_{s + d \cdot (k-1), d} - B_{s-d, d} \right) \end{aligned}

즉, $A_{s, d}$ 와 $B_{s, d}$ 만 전치리를 해놓는다면 $O(1)$ 에 값을 구할 수 있습니다.

$A_{s, d}$ 와 $B_{s, d}$ 의 전처리는 $O(n\sqrt{n})$ 에 할 수 있습니다. 그래서 전체 시간복잡도는 $O((n+q)\sqrt{n})$ 이 됩니다.

$\sqrt{n}$ 을 기준으로 케이스를 나눈 이유는 만약에 $b$ 를 기준으로 케이스를 나누면 시간복잡도는 $O\left( nb+q \cdot \frac{n}{b} \right)$ 이 됩니다. TLE가 나지 않도록 $b$ 값을 설정하면 되고 저는 그 중에서 $\sqrt{n}$ 을 선택했습니다.