けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 1390 Arithmetic Progressions (ICPC アジア 2018 B) (300 点)

数列 DP AOJ ICPCアジア AOJ-ICPC300点 AOJ-ICPC 等差数列

面白かった

問題へのリンク

問題概要

N 個の整数 A_1, A_2, \dots, A_N が与えられる。この中から最大個数の整数を選んで、それを小さい順に並べたときに等差数列となるようにせよ。

制約

  • 3 \le N \le 5000
  • 0 \le A_i \le 10^{9}

考えたこと

すごく DP っぽい雰囲気の問題。イメージとしては

  • dp[ i ] := i 番目の要素が最後になるような最長等差数列の長さ

という感じの DP になりそう。ただこれだけでは公差がわからないので

  • dp[ i ][ diff ] := i 番目の要素が最後になるような公差が diff の最長等差数列の長さ

という感じにしたい。これでもまだ diff が 109 くらいになってしまうのでダメで

  • dp[ i ][ j ] := i 番目の要素が最後になるような公差が A[ i ] - A[ j ] となる最長等差数列の長さ

としてあげると良さそう。これは言い換えると、最後の 2 要素が A[ j ] と A[ i ] ということになる。このとき、

  • A[ i ] - A[ k ] = A[ j ] - A[ i ]

となるような k がもしあったならば、

  • dp[ i ][ j ] = dp[ k ][ i ] + 1

という風になる。そのような k がなかったら dp[ i ][ j ] = 2 になる。i とか j とか k の順番が混乱しそうだったのでメモ化再帰で書いてみた。計算量は本来は O(N^{2}) なんだろうけど、僕の実装だと O(N^{2} \log{N}) になってしまってそうだけど、解説を見ると、しゃくとり法に近い要領で O(N^{2}) でもできそう。

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;

vector<int> a;
map<int,int> ma;

int dp[5100][5100] = {0};
int rec(int i, int j) {
    if (dp[i][j] != 0) return dp[i][j];
    int val = a[i] * 2 - a[j];
    if (!ma.count(val)) return 2;
    int k = ma[val];
    int res = rec(k, i) + 1;
    return dp[i][j] = res;
}

int main() {
    int N; cin >> N;
    a.resize(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    sort(a.begin(), a.end());
    
    for (int i = 0; i < N; ++i) ma[a[i]] = i;
    int res = 2;
    for (int i = 0; i < N; ++i) {
        for (int j = i + 1; j < N; ++j) {
            res = max(res, rec(i, j));
        }
    }
    cout << res << endl;
}