https://kyopro.hateblo.jp/entry/2018/12/09/221608

問題. Arithmetic Progressions

非負整数からなる集合 $V = \{v_1, \ldots, v_n\}$ が与えられる． $V$ の部分集合で等差数列となるものの中で要素数最大のものを見つけ，その要素数を答えよ．

制約： $2 \le n \le 5000$ ，　 $0 \le v_i \le 10^9$

解法1

初めに， $V$ を昇順にソートした列 $(v'_1, \ldots, v'_n)$ を考える．次に，任意の $1 \le i < j \le n$ に対して，初項 $v'_{i}$ で公差が $d = v'_{j} - v'_{i}$ の等差数列で左と右に極大なものを考える．左に極大とは $v'_{i} - d$ となる $V$ の要素が存在しなことであり，集合 $V$ を std::set で管理すると $O(\log n)$ 時間で判定できる．左が極大なものに対して， $v'_{j}$ から順番に右に公差 $d$ の要素が存在するかを $O(\log n)$ 時間で判定して，存在しなくなるまで繰り返すと右に極大な数列を得ることができる．このような等差数列の中で要素数最大のものが答えとなる．
極大な数列だけを考えているので，任意の $1 \le i < j \le n$ に対して， $v'_{i}, v'_{j}$ の組は高々2回しか見ない．また， $v'_{i}, v'_{j}$ のアクセスは定数時間か $O(\log n)$ 時間しかかからないので，全体で $O(n^2 \log n)$ 時間で答えを求めることができる．

計算時間： $O(n^2 \log n)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n;
    cin >> n;

    vector<int> v(n);
    set<int> memo;
    for (int i = 0; i < n; ++i) {
        cin >> v[i];
        memo.insert(v[i]);
    }
    sort(v.begin(), v.end());

    int max_len = 2;
    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            int diff = v[j] - v[i];
            if (i != 0 && memo.find(v[i] - diff) != memo.end()) continue;

            int cur = v[j] + diff, len = 2;
            while (memo.find(cur) != memo.end()) {
                ++len;
                cur += diff;
            }
            max_len = max(max_len, len);
        }
    }

    cout << max_len << endl;

    return 0;
}

解法2

解法1の高速化を考える．初項 $v'_{i}$ で公差が $d = v'_{j} - v'_{i}$ の右に極大な等差数列を考えていたときに，nxt[i][j] = k ( $v'_{k} = v'_{j} + d$ ) という二次元配列があれば十分であることが分かる．なぜならば， $v'$ の任意の等差部分数列 $(v'_{i_1}, v'_{i_2}, \ldots, v'_{i_m})$ と，任意の $1 \le j < m - 1$ に対して，nxt[j][j + 1] = j + 2 となるからである．左に極大という条件も， $v'$ の左から順番に初項の候補を探索していたので，nxt[i][j] は一度だけ使用可能とすれば満たされる．したがって，もし nxt[i][j] を $O(n^2)$ 時間で求めることができれば，答えも $O(n^2)$ 時間で求めることができる．実際，任意の $1 \le i < j \le k \le n$ に対して，nxt[i][j] <= nxt[i][k] となり nxt[i] の値は単調増加するので，nxt[i] の値を $O(n)$ 時間で計算可能である．

計算時間： $O(n^2)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n;
    cin >> n;

    vector<int> v(n);
    for (int i = 0; i < n; ++i) cin >> v[i];
    sort(v.begin(), v.end());

    vector<vector<int>> nxt(n, vector<int>(n, -1));
    for (int i = 0; i < n; ++i) {
        int idx = i;
        for (int j = i + 1; j < n; ++j) {
            const int diff = v[j] - v[i];
            while (idx < n && v[idx] < v[j] + diff) ++idx;
            if (idx < n && v[idx] == v[j] + diff) nxt[i][j] = idx;
        }
    }

    int max_len = 2;
    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            if (nxt[i][j] == -1) continue;

            int prev = i, cur = j, len = 2;
            while (nxt[prev][cur] != -1) {
                int tmp_prev = prev, tmp_cur = cur;
                cur = nxt[tmp_prev][tmp_cur];
                prev = tmp_cur;
                nxt[tmp_prev][tmp_cur] = -1;
                ++len;
            }
            max_len = max(max_len, len);
        }
    }

    cout << max_len << endl;

    return 0;
}

まとめ

計算時間の見積もりが少し難しい