solved.ac Grand Arena #3

풀이 1

사용할 과일을 두 종류 고릅니다. 그 이후, 두 종류의 과일로 이루어진 가장 긴 연속된 부분수열을 구합니다. 가장 긴 연속된 부분수열을 구하는 방법은 언어에서 지원하는 split 함수 등을 사용하거나, cnt변수를 관리하며 앞쪽부터 훑어보는 방법이 있습니다. 원하는 종류의 과일일 때는 cnt를 $1$ 더하고, 원하지 않는 종류일 경우에는 cnt를 $0$으로 초기화 해주면서, cnt의 최댓값을 구하면 됩니다.

과일의 종류 $d = 9$라고 하면, 과일을 두 개 고르는 방법이 $\frac{d(d-1)}{2} = 36$가지 있고, 이 각각에 대해서 $\mathcal{O}(N)$ 시간이 듭니다. 총 시간복잡도는 $\mathcal{O}(d^2N)$입니다.

코드 (C++20)

#include <iostream>
#include <ranges>
#include <vector>
using namespace std;

int main()
{
    int N;
    cin >> N;
    vector<int> S(N);
    for (int &s : S)
        cin >> s;
    int ans = 0;
    for (int i = 1; i <= 9; i++)
        for (int j = 1; j < i; j++)
        {
            auto v = S | views::transform([&](int x)
                                          { return x == i || x == j; });
            for (auto w : views::split(v, false))
                ans = max(ans, (int)distance(w.begin(), w.end()));
        }
    cout << ans << endl;
}

코드 (Python)

import re


def main():
    input()
    s=''.join(input().split())
    ans = 0
    for i in range(1, 10):
        for j in range(1, i):
            ans = max(ans, max(map(len, re.split(f'[^{i}{j}]', s))))
    print(ans)


if __name__ == "__main__":
    main()

코드 (Rust)

use std::io::*;
use std::iter::FromIterator;

fn main() {
    stdin().read_line(&mut String::new()).unwrap();
    let mut s = String::new();
    stdin().read_line(&mut s).unwrap();
    let s = String::from_iter(s.chars().filter(|&x| !x.is_ascii_whitespace()));
    let mut ans = 0;
    for i in '1'..='9' {
        for j in '1'..i {
            ans = ans.max(
                s.split(|x| x != i && x != j)
                    .map(|x| x.len())
                    .max()
                    .unwrap_or_default(),
            );
        }
    }
    println!("{}", ans);
}

풀이 2

$S_l$로 시작해서, 두 종류 이하의 과일이 등장하는 최대 길이 구간의 오른쪽 끝 점을 $R_l$이라고 합시다. 즉, $S_l, S_{l+1}, \cdots, S_{R_l}$에서 두 종류 이하의 과일이 등장합니다. $S_l$을 제외한 $S_{l+1}, \cdots, S_{R_l}$도 두 종류 이하의 과일이 등장하기 때문에, $R_{l+1}$은 $R_l$보다 항상 크거나 같습니다. 즉, $l$에 따라 $R_l$은 증가하는 형태를 보이게 됩니다.

$R_{l-1}$이 계산되었다고 할 때, $R_l$이 $R_l$보다 크다는 것을 알 기 때문에, $r = R_{l-1}$로 시작해서 $r$을 $1$씩 증가시켜보면서 $S_{l}, \cdots, S_r$에서 두 종류 이하의 과일이 등장하는 최대 $r$을 naive하게 반복문으로 구현해도 됩니다. 반복문을 돌리면서 $r$은 총 $N$ 증가하기 때문에, 확인하는 총 횟수는 $O(N)$번이 넘지 않습니다.

이제 $S_l, \cdots, S_r$에 있는 과일 종류를 빠르게 구하면 됩니다. 이는 길이 $10$의 cnt 배열을 관리해, 각 과일이 현재 보고 있는 $S_l, \cdots, S_r$ 구간에 몇 개 있는지를 저장하고, 동시에 $1$개 이상 존재하는 과일이 몇 종류인지를 저장하면 $O(1)$에 문제를 풀 수 있습니다. 총 시간복잡도는 $O(N)$입니다.

아래 코드에서 변수 r 은 $R_l + 1$을 관리하도록 구현되어있습니다.

코드 (C++)

#include <bits/stdc++.h>
using namespace std;
int main()
{
    int N;
    cin >> N;
    vector<int> S(N);
    for (int &s : S)
        cin >> s;
    int ans = 0, r = 0, non_zero = 0;
    vector<int> cnt(10);
    for (int l = 0; l < N; l++)
    {
        while (r < N)
        {
            int nxt_non_zero = non_zero + (cnt[S[r]] == 0 ? 1 : 0);
            if (nxt_non_zero <= 2)
            {
                cnt[S[r]]++;
                r++;
                non_zero = nxt_non_zero;
            }
            else
                break;
        }
        ans = max(ans, r - l);
        cnt[S[l]]--;
        if (cnt[S[l]] == 0)
            non_zero--;
    }
    cout << ans << endl;
}

코드 (Python)

def main():
    N = int(input())
    S = map(int, input().split())

    ans = r = non_zero = 0
    cnt = [0]*10

    for l in range(N):
        while r < N:
            nxt_non_zero = non_zero + (1 if cnt[S[r]] == 0 else 0)
            if nxt_non_zero <= 2:
                cnt[S[r]] += 1
                r += 1
                non_zero = nxt_non_zero
            else:
                break
        ans = max(ans, r - l)
        cnt[S[l]] -= 1
        if cnt[S[l]] == 0:
            non_zero -= 1

    print(ans)


if __name__ == "__main__":
    main()

코드 (Rust)

use std::io::*;

fn main() {
    let n = read_vec()[0];
    let s = read_vec();

    let mut ans = 0;
    let mut r = 0;
    let mut non_zero = 0;
    let mut cnt = [0; 10];

    for l in 0..n {
        while r < n {
            let nxt_non_zero = non_zero + if cnt[s[r] as usize] == 0 { 1 } else { 0 };
            if nxt_non_zero <= 2 {
                cnt[s[r]] += 1;
                r += 1;
                non_zero = nxt_non_zero;
            } else {
                break;
            }
        }
        ans = ans.max(r - l);
        cnt[s[l]] -= 1;
        if cnt[s[l]] == 0 {
            non_zero -= 1;
        }
    }

    println!("{}", ans);
}

fn read_vec() -> Vec<usize> {
    let mut s = String::new();
    stdin().read_line(&mut s).unwrap();
    s.split_whitespace().map(|x| x.parse().unwrap()).collect()
}