BOJ 1994 등차수열

문제 링크입니다.

문제 요약

주어진 $N$개의 수로 최대한 긴 등차수열을 만드세요.

풀이

주어진 수들을 오름차순 정렬합니다. 원하는 값을 ${\cal O}(\lg N)$에 찾을 수 있게 하기 위함입니다. 공차는 음이 아닌 정수로 제한하겠습니다. 공차가 음수인 경우는 탐색 방향만 다른 거라 굳이 해 볼 필요 없습니다.

등차수열 상 연속한 두 수의 쌍으로 가능한 것은 ${\cal O}(N^2)$가지입니다. 두 값을 각각 $a \le b$라고 하면, 등차수열 상 다음에 올 수 있는 값은 $2b-a$로 고정됩니다. 해당 값이 존재하는 지를 이분 탐색으로 찾으면 ${\cal O}(N^2 \lg N)$에 모든 쌍에 대한 다음 수의 위치를 찾아낼 수 있습니다.

이를 이용해 $d[i][j]$를 첫 두 수가 $A_i, A_j$인 등차수열의 최대 길이로 정의하고 DP를 수행하면 우리가 원하는 답을 얻어낼 수 있습니다.

코드

#include<bits/stdc++.h>
using namespace std;
map<int, int> cnt;
vector<int> v;
int d[2020][2020];

int dp(int i, int j){
	int& ret = d[i][j];
	if(ret) return ret;
	ret = 2;
	int nxt = lower_bound(v.begin()+j, v.end(), 2*v[j]-v[i]) - v.begin();
	if(v[nxt] != 2*v[j]-v[i]) return ret;
	return ret = 1 + dp(j, nxt);
}

int main(){
	cin.tie(0);cout.tie(0);ios::sync_with_stdio(0);
	int n, ans = 2; cin >> n;
	if(n == 1) return !(cout << 1);
	for(int i = 0; i < n; i++){
		int k; cin >> k;
		if(cnt.find(k) != cnt.end()){
			cnt[k]++;
			ans = max(ans, cnt[k]);
		}
		else{
			cnt.insert({k,1});
			v.push_back(k);
		}
	}
	sort(v.begin(), v.end());
	for(int i = 0; i < v.size(); i++){
		for(int j = i+1; j < v.size(); j++){
			ans = max(ans, dp(i, j));
		}
	}
	cout << ans;
}