[백준 BOJ 알고리즘] 수열에서 a_j - a_i = a_k - a_j인 i < j < k 구하기

길이가 N인 수열에서 등차수열 관계를 가지는 세 수를 뽑는 경우의 수는 어떻게 구할 수 있을까요?

당연히 N과 원소의 크기 a_i의 범위가 모두 크다면 풀이할 수 없겠지만, 특정 제한이 있다면 풀이할 수 있습니다.

이 포스트에서는 세 가지의 경우에 따라 풀이 방법을 다르게 하여 푸는 방법을 소개합니다.

(참고로 이 문제들은 제가 나중에 따로 볼 일이 있을 것 같아 포스트를 분리하여 작성합니다.)

 

 

백준 BOJ 13558번 : 등차수열의 개수

문제 난이도 : Gold III

알고리즘 분류 : 브루트포스 (와 비슷한..)

 

길이가 10만 이하의 N이고 각 원소가 1 ~ 30,000 사이인 수열이 있을 때 a_j - a_i = a_k - a_j이고 i < j < k를 만족하는 순서쌍의 수를 구하는 문제입니다.

 

공략해야 하는 핵심은 원소의 크기가 1 ~ 30,000으로 제한되어 있다는 것입니다.

따라서 원소를 하나씩 입력받으면서, 이 원소가 들어왔을 때 새롭게 발생하는 등차수열의 개수를 더해줍니다.

그 다음 이 원소가 앞의 원소와 구성하는 쌍들을 찾아 뒤에 어떤 수가 나온다면 등차수열이 몇 개가 발생하는지를 저장해줍니다.

이 과정을 원소를 N개 입력받는동안 N번 반복해주면 됩니다.

 

시간 복잡도는 수열의 크기 제한이 30,000이므로 대략 O(30,000 N)입니다.

N이 10만 이하이므로 최대 연산 횟수를 30억회 정도로 볼 수 있는데, 그럼에도 3초라는 제한 시간 안에 통과할 수 있습니다.

 

 

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

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL), cout.tie(NULL);

    int N; cin >> N;

    vector<int> v(30001), u(30001);

    int ans = 0;
    while(N--) {
        int x; cin >> x;

        ans += u[x] + v[x] * (v[x] - 1) / 2;

        v[x]++;

        for(int i=1; i<x; i++)
            if(x + (x - i) <= 30000) u[x + (x - i)] += v[i];
        for(int i=30000; i>x; i--)
            if(x - (i - x) >= 1) u[x - (i - x)] += v[i];
    }

    cout << ans << "\n";
}

 

 

백준 BOJ 18079번 : Managing Difficulties

문제 난이도 : Silver II

알고리즘 분류 : 해시를 사용한 집합과 맵

 

이 문제 역시 길이 N인 수열에서 a_j - a_i = a_k - a_j이며 i < j < k인 (i, j, k)의 순서쌍의 개수를 구하는 문제입니다.

다만 이 문제는 원소 a_i의 범위가 1 ~ 10^9입니다.

대신 이 문제에서는 N의 범위가 2,000이하로 작아 O(N^2)의 풀이를 생각해볼 수 있습니다.

 

식 a_j - a_i = a_k - a_j에서 a_i만 넘기고 넘겨주면 a_i = 2 * a_j - a_k이므로, a_i를 한 칸씩 갱신하면서 j, k에 대해 2 * a_j - a_k = a_i를 만족하는 조합들을 찾아주면 됩니다.

이 때 a_i 값의 저장은 map을 통해서 수행해줍니다.

 

이 때 그냥 map을 사용하면 map의 작동 속도가 느리므로 시간 초과에 걸리게 됩니다. (map의 작동 시간은 대략 log 시간으로 보면 된다고 합니다.)

중복이 없는 1대1 매칭이라면 속도가 더 빠른 unordered_map을 사용할 수 있으므로 이를 이용하여 풀이해주어야 정답 처리를 받을 수 있습니다.

 

 

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

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL), cout.tie(NULL);

    int T; cin >> T;

    while(T--) {
        int N; cin >> N;

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

        unordered_map<int, int> m;
        int ans = 0;

        for(int j=1; j<N-1; j++) {
            int i = j-1;
            m[v[i]]++;

            for(int k=j+1; k<N; k++) ans += m[v[j]*2 - v[k]];
        }

        cout << ans << "\n";
    }
}

 

 

백준 BOJ 13423번 : Three Dots

문제 난이도 : Silver II

알고리즘 분류 : 정렬, 이분 탐색

 

여러 테스트케이스마다 N개의 x축 상의 좌표가 주어질 때, 등차수열 관계인 세 점의 순서쌍의 수를 구하는 문제입니다.

 

이 문제는 N도 1,000이하이지만 테스트케이스의 수가 많아 시간 초과가 발생하는 경우입니다.

다행히 이 문제는 수를 정렬해도 답에 영향을 주지 않으므로 정렬한 이후 i, j에 대한 k의 이분탐색을 통해 O(T N^2 log N)에 풀이가 가능합니다.

 

 

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

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL), cout.tie(NULL);

    int T; cin >> T;

    while(T--) {
        int N; cin >> N;

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

        sort(v.begin(), v.end());

        int ans = 0;
        for(int i=0; i<N; i++)
            for(int j=i+1; j<N; j++) {
                ans += upper_bound(v.begin()+j+1, v.end(), v[j]*2 - v[i])
                       - lower_bound(v.begin()+j+1, v.end(), v[j]*2 - v[i]);
            }

        cout << ans << "\n";
    }
}