[ 소프티어 ] 자동차 테스트 (C++)

https://softeer.ai/practice/6247

 

 3개의 수를 골라 중앙값이 m이 되도록 하는 경우는 m을 고르고 m보다 작은 수 하나, m보다 큰 수 하나를 고르는 경우와 같다. 

 따라서 m이 주어졌을 때, (m보다 작은 수 중 하나를 고르는 경우의 수) x (m보다 큰 수 중 하나를 고르는 경우의 수)를 출력하면 된다.

 

이 때, m보다 작은 수(or 큰 수) 중 하나를 고르는 경우의 수는 m보다 작은 수(or 큰 수)의 개수를 k라고 했을 때 _kC₁ = k와 같다.

따라서 (m보다 작은 수의 개수) x (m보다 큰 수의 개수)를 출력하여 문제를 해결할 수 있다. 최대 200000번 반복해야하기 때문에 m보다 작은 수와 큰 수의 개수를 구하는 알고리즘은 O(log n)의 시간복잡도를 가져야 한다.

 

[ 풀이 방법 ]

1. 자동차의 연비 n개를 정렬한다.

  - STL sort(): O(n log n)

 

2. m을 입력받고, (m보다 작은 수의 개수 x m보다 큰 수의 개수)를 구한다.

  - STL lower_bound()와 upper_bound()를 이용하면 m보다 작은 수와 큰 수의 개수를 O(log n) 시간안에 구할 수 있다.

  - 입력된 자동차 연비 중 m이 존재하지 않는다면 0을 출력한다.

 

3. 위 과정을 q번 반복한다.

 

[ 회고 ]

 입력된 자동차 연비 중 m이 존재하는지 찾기 위해 STL find()를 사용하였으나 O(n)의 시간복잡도를 갖기 때문에 시간 초과가 발생했다.

lower_bound()를 이용해 입력된 자동차 연비 중 m이 존재하는지 찾도록 수정하여 정답을 맞출 수 있었다.

 

STL find()의 시간복잡도를 착각하고 있었는데 바로잡을 수 있었다.

 

[ 정답 코드 ]

#include<iostream>
#include <algorithm>
#include <vector>

using namespace std;

typedef long long ll;

int main(int argc, char** argv)
{
    iostream::sync_with_stdio(0);
    cout.tie(0); cin.tie(0);

    ll n, q;
    cin >> n >> q;
    
    vector<ll> car(n);
    for(int i = 0; i < n; i++) {
        cin >> car[i];
    }
    sort(car.begin(), car.end());

    ll t;
    while(cin >> t) {
        if(lower_bound(car.begin(), car.end(), t) == car.end() || car[lower_bound(car.begin(), car.end(), t) - car.begin()] != t) {
            cout << 0 << '\n';
            continue;
        }
        ll left = lower_bound(car.begin(), car.end(), t) - car.begin();
        ll right = car.size() - (upper_bound(car.begin(), car.end(), t) - car.begin());
        cout << left * right << '\n';
    }
    
   return 0;
}

 

혹은 m보다 작은 수와 큰 수의 개수를 구하기 위해 map을 사용하여 해결할 수 있다.

#include<iostream>
#include <algorithm>
#include <vector>
#include <map>

using namespace std;

typedef long long ll;

int main(int argc, char** argv)
{
    iostream::sync_with_stdio(0);
    cout.tie(0); cin.tie(0);

    ll n, q;
    cin >> n >> q;
    
    vector<ll> car(n);
    map<ll, int> left;
    map<ll, int> right;
    for(int i = 0; i < n; i++) {
        cin >> car[i];
    }
    sort(car.begin(), car.end());
    for(int i = 0; i < car.size(); i++) {
        left[car[i]] = i;
        right[car[i]] = i+1;
    }

    ll t;
    while(cin >> t) {
        if(!left[t]) {
            cout << 0 << '\n';
            continue;
        }
        ll l = left[t];
        ll r = car.size() - right[t];
        cout << l * r << '\n';
    }
    
   return 0;
}