BOJ 13558 등차수열의 개수

난이도 : ???

 

이 문제 같은 경우에는, 굉장히 어려운 (그리고 재밌는) 정해와, 그럼에도 불구하고 솔브드 난이도가 골드 3이 되어버린 이상한 스토리가 있다. 그래서 난이도는 ??? 인데, 정해를 기준으로 생각한다면 최소 다이아 급의 문제이고, 그냥 뚫을 생각이라면 골드5? 골드5정도면 충분한것 같다.

---

1. 문제 설명

 

길이가 $N$ 인 수열 $A_1, \dots A_n$ 에 대하여, $A_i, A_j, A_k$ 가 등차수열이 되는 $(i, j, k)$ (순서 지켜야 한다) 순서쌍의 개수를 찾는 문제이다. 수열의 길이는 10만, 각 수는 3만 이하.

 

2. 풀이

2-1. Multiset에 대해 생각하기

먼저, 이 문제가 왜 망했는지(...) 의 이야기를 하기 전에, 복잡도가 실제로 낮은 풀이에 대해 생각해 보자. 정해는 코딩이 미친듯이 어려워 보여서 코딩을 할 엄두도 내지 못했다. 이 풀이는 팀원인 diordhd와 갓-갓 tlwpdus 가 각자 따로 생각해서 알려줬고, 대충 나도 생각 방향은 비슷했지만 최종 스텝을 마치지 못했었다.

 

각 수가 3만 이하라는 사실에 주목해 보자. 각 수가 3만 정도로 작다는 것은 일단 당연히 카운팅을 하고 싶어 보인다.

그런데, 수열에서 인덱스의 순서를 지켜야 하는 상황에서는 도저히 카운팅을 할 수가 없다. 즉, "12,723" 이 "12번" 나왔다는 정보만 가지고는 아무것도 할 수 없고, 이 수들이 실제로 어디에 나왔는지를 확실히 알아야 한다.

 

이 문제를 조금 변형해서, 인덱스 순서를 무시해도 된다고 생각해 보자. 즉, 주어진 문제를 sequence 가 아닌, multiset 형태에 대하여 푸는 경우이다. multiset에서는 당연히 일단 카운팅을 한 다음, 카운팅한 배열 $\texttt{CT}$ 을 들고 있도록 하자. 

 

어떤 세 수 $a, b, c$ 가 저 순서대로 등차수열일 필요충분조건은 $a+c = 2b$ 이다. 즉, 각각의 수 $a_i$ 들에 대하여, 합이 $2a_i$ 가 되는 $(j, k)$ 들의 순서쌍의 개수를 찾는 문제로 환원할 수 있고, 이 문제는 $O(nk)$ 에 풀 수 있다.

중요한 사실 하나는, 이  Naive 풀이가 수열에서도 된다는 것이다. (별로 어렵지 않게 할 수 있다) 조금만 설명을 하자면, 어떤 원소 $a_j$ 입장에서 왼쪽까지의 카운터와 오른쪽까지의 카운터 배열을 각각 두고, 그 두개에서 서로 맞는 경우를 찾으면 한 원소를 가운데 원소로 하는 triplet을 $O(k)$ 에 셀 수 있다. 자명하게, 원한다면 카운터 없이 그냥 직접 돌아서 $O(n^2)$ 에 구할 수도 있다. 

 

BOJ에도 이 비슷한 문제가 여러 개 있는데 (당장 생각나는 문제는 https://www.acmicpc.net/problem/17134), 이 경우 CT 배열을 다항식으로 간주하고 자기 자신과 다항식곱셈 (Convolution) 을 수행할 수 있다. 예를 들어, 4가 한 번, 3이 세 번 있다면 

$$ct(x) = x^4 + 3x^3$$ 이 되고, 자기 자신과의 곱은 $$f(x)^2 = x^8+6x^7+9x^6$$ 가 된다. 즉, 이 배열에서 두 개를 뽑아서 8을 만들 수 있는 경우가 1가지, 7을 만들 수 있는 경우가 6가지, 6을 만들 수 있는 경우가 9가지인데, 이때 자기 자신과의 곱은 두번 세어졌으므로 $x^8$ 이나 $x^6$은 각각 한번, 세번씩 빼야 하고, 뽑는 순서를 고려하지 않으므로 실제로는 7을 만드는 경우가 3가지, 6을 만드는 경우가 3가지이다.

 

즉, 다항식 곱셈을 하면 "각각의 수 $a_i$ 들에 대하여, 합이 $2a_i$ 가 되는 $(j, k)$ 들의 순서를 찾는 문제" 를 $O(k^2)$ 에 풀 수 있다는 것이다. ($k = 30,000$) 여기까지만 해도 사실 Platinum 정도 난이도의 FFT 기본 문제인 것 같다.

 

 

2-2. 수열에 대해 생각하기

수열에 대한 아이디어를 떠올리기 위해, 저 풀이를 어떻게 수열에 적용할 수 있을지 생각해 보자. 먼저, 만약 다항식을 곱셈하여 이 문제를 해결하고 싶다면, 다음과 같이 다항식 곱셈을 쓰면 된다 .

$L(x, i)$ : $i$ 번보다 작은 인덱스까지만 센 CT 배열을 다항식으로 표현

$R(x, i)$ : $i$ 번 보다 큰 인덱스까지만 센 CT 배열을 다항식으로 표현

이제, $L(x, i)R(x. i)$ 를 FFT 곱셈하고, 그 곱셈 결과에서 $2a_i$ 차 항의 계수를 세면 된다. 이 방법은 $n$ 번의 FFT를 요구하며, $O(k \log k)$ 에 다항식 곱셈이 돌아가므로 수행 시간은 $O(nk \log k)$ 이 된다.  

 

이렇게는 어림도 없지만, 여기서 중요한 관찰은 인덱스 하나를 지나갈 때, 그렇게 많은 값이 바뀌지 않는다는 점이다. $L(x, i)$ 와 $L(x, i+1)$ 의 차이는 $a_i$ 차 항에서만 달라지고, $R$ 도 한 개만 달라진다. 굳이 그러면 그렇게 많은 FFT를 할 필요가 없어 보인다. 

 

전체를 $\sqrt{n}$ 개 정도의 버킷으로 나눈다고 생각해 보자. 그러면 지금 내가 보는 위치의 왼쪽에 있는 버킷들에 대한 Count 배열을 관리하는것은 어렵지 않고, 같은 방법으로 오른쪽에 있는 버킷들에 대한 Count 배열도 관리가 어렵지 않다.

오른쪽 버킷과 왼쪽 버킷에 대한 Count 를 유지하고 있으므로, 내가 보는 원소(나)를 가운데로 각 원소의 입장에서 Arithmetic Sequence triplet 을 다음과 같이 나누어 생각할 수 있다.

• 왼쪽은 나보다 왼쪽 버킷에 있고, 오른쪽도 나보다 오른쪽 버킷에 있는 경우
• 왼쪽은 나보다 왼쪽 버킷에 있고, 오른쪽은 나랑 같은 버킷에 있는 경우
• 왼쪽은 나랑 같은 버킷에 있고, 오른쪽은 나보다 오른쪽 버킷에 있는 경우
• 양쪽 원소가 모두 나랑 같은 버킷에 있는 경우

각 경우를 생각해 보자.

(1) 의 경우, 이제 왼쪽 버킷에 대한 $L(x)$ 와 $R(x)$를 $O(k \log k)$ 시간에 곱하면, 버킷 전체의 (1) 번 경우를 한번에 구할 수 있다. (직접 세는 과정은 FFT보다 빠르니까 복잡도 계산에 영향 X)

(2)와 (3) 의 경우, 왼쪽 / 오른쪽에서부터 직접 세면서, 가지고 있었던 오른쪽 카운터를 이용해서 한번 훑어보면 구할 수 있다. 버킷 사이즈 $B$ 에 대해 $O(B)$ 고, 이는 $O(\sqrt n)$.

(4) 의 경우, 나이브하게 위에서 설명한 솔루션으로 세자. 그러나 이경우 한 버킷에 원소 개수가 $\sqrt{n}$ 개씩이므로, $O(\sqrt{n}^2)$ 가 더 이득이다. 

그러면, 총 시간 복잡도는 $O(k \log k + n)$ 이고, 이것을 버킷마다 해야 하므로, $O(\sqrt{n}k\log{k} + n \sqrt{n})$ 복잡도에 전체 문제를 해결할 수 있다.

 

 

3. 여담

저런 솔루션에도 불구하고 문제가 골드 3인 이유는, 3초의 시간 제한이 꽤 넉넉하고 현대 컴퓨터의 발전이 눈부시다 보니 (....) $O(nk)$ 솔루션이 $n = 100, 000$, $k = 30,000$ 에 대해서 3초에 돌기 때문이다. 왜 도는지는 잘 모르겠고, 혹시 시간 제한을 줄인다면 다음과 같은 문제가 있다.

• 설명한 정해의 시간 복잡도가 은근히 크고, FFT는 상수가 큰 연산인 데 반해 카운터 배열 가지고 세는건 상수가 정말 작기 때문에, FFT를 쓰면 맞고 안쓰면 틀리는 시간 제한을 잡는 것이 어렵다.
• $O(nk)$ 에 AVX로 상수커팅하면 /8되고, 그러면 슥삭... FFT를 아무리 잘 짜도 (AVX 쓰는 FFT도 힘들어 보이는데 모르겠다) 이거보다 빠르게 만들기는 힘들어 보인다.

 

4. 코드

https://www.acmicpc.net/source/share/b659d90250f04efca0238c2e3330c811