[백준 25571번] 지그재그 부분배열

문제 링크

https://www.acmicpc.net/problem/25571

 

문제 요약

길이가 \(N\)인 배열이 주어졌을 때, 만들 수 있는 연속된 부분배열들 중 지그재그인 부분배열의 개수를 찾는 문제이다. 배열이 지그재그라는 뜻은 배열의 원소가 감소와 증가 또는 증가와 감소를 반복하는 배열을 의미한다.

 

문제 풀이

만들 수 있는 모든 부분배열 \(A[lo, hi]\)에 대해서 지그재그인지 검사하면 답을 구할 수 있다. 하지만 이 방법은 시간복잡도가 \(\mathcal {O}(N^2)\)이라서, \(N\)이 10만 일 때 제한시간 내에 문제를 해결하지 못하므로 더 효율적인 풀이가 필요하다.

 

지그재그 부분배열의 특징을 살펴보면 \(A[lo, hi - 1]\)이 지그재그이고, \(A[lo, hi]\)가 지그재그일 때 \(A[lo ... hi - 1, hi]\) 또한 지그재그 부분배열이다. 즉 \(A[lo, hi - 1]\)이 지그재그일 때 \(A[lo, hi]\)도 지그재그이면 지그재그인 부분배열의 개수가 \(hi - lo\) 개만큼 추가되는 것이라고 볼 수 있다. 반대로 \(A[lo, hi - 1]\)이 지그재그인데 \(A[lo, hi]\)가 지그재그가 아니라면, \(A[lo ... hi - 1, hi]\)는 전부 지그재그 부분배열이 될 수 없다. 이 점을 이용하면 만들 수 있는 부분배열을 전부 순회하지 않고도, \(hi\)를 하나씩 늘려가면서 답을 찾을 수 있다.

 

소스 코드

아이디어는 간단하지만 구현은 조금(?) 까다롭다. \(A[lo, hi - 1]\)이 마지막이 감소하는 지그재그인지, 증가하는 지그재그인지에 따라 \(A[hi - 1]\)과 \(A [hi]\)의 증가 또는 감소를 체크해야 한다. 또한 \(A[lo, hi - 1]\)이 지그재그이지만 \(A[lo, hi]\)가 지그재그가 아닐 때, \(A[hi - 1, hi]\)가 또 다른 지그재그 부분배열의 시작이 될 수 있는지도 함께 확인해야 한다.

 

#include <stdio.h>

enum LastState {
    NONE, // 초기 OR 지그재그가 아닌 상태
    DESC, // 마지막이 감소 상태
    INC, // 마지막이 증가 상태
};

int n, t, A[100010];

int main() {
    scanf("%d", &t);
    
    while (t--) {
        scanf("%d", &n);
        
        for (int i = 0 ; i < n ; ++i) {
            scanf("%d", &A[i]);
        }
        
        int lo = 0;
        long long ans = 0;
        LastState lastState = NONE;
        
        for (int hi = 1 ; hi < n ; ++hi) {
            // 초기 상태이거나 지그재그 부분수열이 아닌 경우
            if (lastState == NONE) {
                lo = hi - 1;
                
                // 새로운 지그재그 부분배열이 시작되는지 체크
                if (A[hi] > A[hi - 1]) {
                    lastState = INC;
                } else if (A[hi] < A[hi - 1]) {
                    lastState = DESC;
                }
            } 
            // 이전 지그재그 부분배열이 감소로 끝난경우
            else if (lastState == DESC) {
                // 증가가 된다면 지그재그 부분배열의 길이가 증가
                if (A[hi] > A[hi - 1]) {
                    lastState = INC;
                }
                // 감소가 된다면 새로운 지그재그 부분배열이 시작
                else if (A[hi] < A[hi - 1]){
                    lastState = DESC;
                    lo = hi - 1;
                }
                // 지그재그 부분배열이 아닌 경우
                else {
                    lastState = NONE;
                }
            }
            // 이전 지그재그 부분배열이 증가로 끝난 경우
            else if (lastState == INC) {
                // 감소가 된다면 지그재그 부분배열의 길이가 증가
                if (A[hi] < A[hi - 1]) {
                    lastState = DESC;
                }
                // 증가가 된다면 새로운 지그재그 부분배열이 시작
                else if (A[hi] > A[hi - 1]) {
                    lastState = INC;
                    lo = hi - 1;
                }
                // 지그재그 부분배열이 아닌 경우
                else {
                    lastState = NONE;
                }
            }
            
            // 지그재그 부분배열이라면 만들 수 있는 모든 부분 배열을 답에 누적
            if (lastState != NONE) {
                ans += hi - lo;
            }
        }
        
        printf("%lld\n", ans);
    }
}