C++에서 부분 리스트를 제거해 k보다 작은 수와 큰 수의 개수를 같게 만드는 프로그램

숫자 리스트 nums와 정수 k가 주어졌다고 가정해 봅시다. 우리는 리스트에서 임의의 연속된 부분 리스트(sublist)를 최대 한 번 제거할 수 있으며, 그 결과 k보다 작은 숫자의 개수와 k보다 큰 숫자의 개수가 서로 같아지는 리스트 중 가장 길이가 긴 것의 길이를 구해야 합니다.

예를 들어, 입력이 nums = [6, 10, 8, 9, 3, 5], k = 6이라면 출력은 5입니다. 부분 리스트 [9]를 제거하면 [6, 10, 8, 3, 5]라는 리스트를 얻을 수 있는데, 이 리스트에는 6보다 작은 수가 [3, 5]로 두 개, 6보다 큰 수가 [10, 8]로 두 개 존재하여 개수가 동일하기 때문입니다.

접근 방식

이 문제는 누적 균형 값과 맵(map)을 활용하면 선형 시간 O(n)에 해결할 수 있습니다. 먼저 각 위치까지의 균형 값을 계산한 뒤, 최종 균형 값을 0으로 만들기 위해 제거해야 하는 가장 짧은 구간을 맵을 통해 찾아내는 방식입니다.

구체적인 단계는 다음과 같습니다.

nums의 크기 + 1만큼의 배열 v를 선언하고 0으로 초기화합니다.
cnt := 0으로 설정합니다.
i := 0부터 nums의 크기 미만까지 반복하며 다음을 수행합니다.
- nums[i]가 k보다 작으면 cnt를 1 증가시킵니다.
- 그렇지 않고 nums[i]가 k보다 크면 cnt를 1 감소시킵니다.
- v[i + 1] = cnt로 저장합니다.
v의 마지막 원소가 0이면 이미 균형이 맞으므로 nums의 크기를 그대로 반환합니다.
delta := v의 마지막 원소로 설정합니다.
맵 m을 하나 선언합니다.
ans := 무한대(INT_MAX)로 초기화합니다.
i := 1부터 v의 크기 이하까지 반복하며 다음을 수행합니다.
- m[v[i] - v의 마지막 원소]의 값이 0이 아니거나, v[i] - v의 마지막 원소가 0이면 ans := min(ans, i - m[v[i] - v의 마지막 원소])로 갱신합니다.
- m[v[i]] := i로 저장합니다.
ans가 여전히 무한대라면 0을 반환하고, 그렇지 않으면 nums의 크기에서 ans를 뺀 값을 반환합니다.

동작 원리

배열 v의 각 원소는 해당 위치까지 살펴봤을 때 '(k보다 작은 수의 개수) - (k보다 큰 수의 개수)'를 의미합니다. 전체 리스트의 최종 균형 값이 0이 아니라면, 특정 구간을 제거해서 이 값을 0으로 만들어야 합니다. 인덱스 j 다음부터 i까지의 구간을 제거하면 최종 균형 값은 v.back() - (v[i] - v[j])가 되므로, v[j] = v[i] - v.back()을 만족하는 지점을 찾으면 됩니다. 맵에 각 균형 값이 처음 등장한 위치를 저장해 두면 제거해야 할 가장 짧은 구간의 길이를 효율적으로 구할 수 있고, 전체 길이에서 그 값을 빼면 정답이 됩니다.

다음 구현 예제를 보면 더 잘 이해할 수 있습니다.

예제

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int solve(vector<int>& nums, int k) {
        vector<int> v(nums.size() + 1, 0);
        int cnt = 0;
        for (int i = 0; i < nums.size(); ++i) {
            if (nums[i] < k)
                ++cnt;
            else if (nums[i] > k)
                --cnt;
            v[i + 1] = cnt;
        }
        if (v.back() == 0) return int(nums.size());
        int delta = v.back();
        map<int, int> m;
        int ans = INT_MAX;
        for (int i = 1; i <= v.size(); ++i) {
            if (m[v[i] - v.back()] != 0 || v[i] - v.back() == 0) {
                ans = min(ans, i - m[v[i] - v.back()]);
            }
            m[v[i]] = i;
        }
        if (ans == INT_MAX)
            return 0;
        else
            return int(nums.size() - ans);
    }
};
main(){
    Solution ob;
    vector<int> v = {6, 10, 8, 9, 3, 5};
    int k = 6;
    cout << ob.solve(v, k); }

입력

{6, 10, 8, 9, 3, 5}, 6