BOJ 16740 - Arithmetic Progressions 링크
(2024.04.24 기준 G1)

문제

길이 $n$의 음이 아닌 정수로 이루어진 수열 $a$가 주어진다. 이 수열에서 일부 숫자를 선택하여 만들 수 있는 가장 긴 등차수열을 만들었을 때, 길이 출력

알고리즘

모든 두 항에 대해 검사한다. 즉 브루트포스

풀이

$n$이 최대 $5\,000$이며, $a$의 원소는 서로 다르다. 또한 시간 제한이 5초로 매우 넉넉하다. 그러므로 배열 정렬 뒤에 모든 $i$, $j$ ($i < j$) 쌍에 대해서 $A_i$와 $A_j$를 각각 등차수열의 1번째, 2번째 항으로 가정하자. 그러면 이 등차수열의 공차 $d$는 $A_j - A_i$가 되고, $n (n > 2)$번째 항은 $A_j + (n-2) \times d$가 된다. 이렇게 $3$, $4$, $\dots$번째 항을 순서대로 찾아가면서 못찾을 때까지 반복하면 된다. 이때, 찾는 방법은 이분 탐색으로 사용하면 $\log n$으로 줄일 수 있다.

코드

• C++

#include <bits/stdc++.h>
using namespace std;

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

    int N; cin >> N;
    int A[N]; for (int i = 0; i < N; i++) cin >> A[i];

    // 배열 정렬 뒤에 가능한 모든 i, j (i < j) 쌍에 대해서
    // Ai와 Aj를 각각 1번째, 2번째 항으로 가정해서
    // 셋째, 넷째 항을 이분 탐색으로 찾아나가면 된다.
    sort(A, A + N);
    int ans = 0;
    for (int i = 0; i < N - 1; i++) for (int j = i + 1; j < N; j++){
        int d = A[j] - A[i]; // 공차
        int k = j; // 찾은 마지막 항의 인덱스
        int l = k + 1, r = N - 1; // 다음 항을 찾을 범위
        int res = 2;
        while (l <= r){
            int st = l, en = r;
            while (st < en){
                int mid = (st + en) >> 1;
                if (A[mid] < A[k] + d) st = mid + 1;
                else en = mid;
            }
            if (A[st] != A[k] + d) break; // 찾지 못했으면 중단
            ++res; k = st; l = k + 1;
        }
        ans = max(ans, res);
    }

    cout << ans;
}

• Python (PyPy3)

import sys; input = sys.stdin.readline

N = int(input())
A = sorted(map(int, input().split()))

# 배열 정렬 뒤에 가능한 모든 i, j (i < j) 쌍에 대해서
# Ai와 Aj를 각각 1번째, 2번째 항으로 가정해서
# 셋째, 넷째 항을 이분 탐색으로 찾아나가면 된다.
ans = 0
for i in range(N - 1):
    for j in range(i + 1, N):
        d = A[j] - A[i] # 공차
        k = j # 찾은 마지막 항의 인덱스
        l = k + 1; r = N - 1 # 다음 항을 찾을 범위
        res = 2
        while l <= r:
            st = l; en = r
            while st < en:
                mid = (st + en) >> 1
                if A[mid] < A[k] + d:
                    st = mid + 1
                else:
                    en = mid
            if A[st] != A[k] + d: # 찾지 못했으면 중단
                break
            res += 1
            k = st
            l = k + 1
        ans = max(ans, res)

print(ans)