알고리즘) 최대 길이의 연속부분수열

 0과 1로 구성된 길이가 N인 수열이 주어지면, 이 수열에서 최대 k번 0을 1로 변경할 수 있습니다. 최대 k번의 변경을 통해 이 수열에서 1로만 구성된 최대 길이의 연속부분수열을 찾는 프로그램을 작성하는 문제이다. 첫 줄에 수열의 길이인 자연수 N(5<=N<=100000)와 K가 입력되고, 다음 줄에는 수열이 입력된다. 이를 바탕으로 도출되는 최대 길이를 출력하면 된다. (처음에는 최대 길이의 연속부분수열을 출력하는 줄 알았다.)

 

초기 작성 코드

import java.util.*;
  
public class Main {
    private int solution(int n, int k, int[] arr) {
        int answer = 0;
		
		int lt = 0, cnt = 0, max = 0, chg = 0;
		for (int rt = 0; rt < n; rt++) {
			if(arr[rt] == 1) {
				cnt++;
			} else {
				chg++;
				if (chg <= k) {
					cnt++;
				} else {
					if (answer < cnt) {
						answer = cnt;
					}
				
					while(chg > k) {
						cnt--;
						lt++;
						if (arr[lt] == 0) {
							chg--;
							lt++;
						}
					}
					
					while(arr[lt] == 0) {
						chg--;
						cnt--;
						lt++;
					}
				}
			}
		}
		
        return answer;
    }

    public static void main(String[] args) {
        Main al = new Main();
        Scanner sc = new Scanner(System.in);

        int n = sc.nextInt();
		int k = sc.nextInt();
		int[] arr = new int[n];
		for (int i = 0; i < n; i++) arr[i] = sc.nextInt();
        
        System.out.println(al.solution(n, k, arr));
    }
}

 이 문제는 못생긴 코드나마 힌트를 얻지 않고 끝까지 내 힘으로 풀어서 Time Limit Exception없이 통과한 문제다..ㅠㅠ 내가 보기에도 부끄러운 못생긴 코드지만 딴에는 이 전날 작성한 연속부분수열 포스트의 개선 코드에 적용한 방법을 적용하려 노력한 결과물이다. 나름대로 정리도 해가며 풀었는데 기억해야할만한 내용이 있을 수 있으니 여기다 적어둔다.

필요한 변수
1. 연속되기 시작하는 인덱스(lt)
2. 현재 탐색하는 인덱스(rt - main loop idx)
3. 0을 1로 변경한 연속부분수열 내부에 0의 개수(chg)
4. 현재 연속되는 길이(cnt) cf. rt - lt + 1해도 될 듯
5. 최대 길이 수열의 길이(answer)

 

유의할 점
1. 연속 1이거나 0을 1로 변경할 수 있는 경우 lt의 변화 없이 rt만 증가한다.
2. lt가 변경되어야 할 때는 'chg <= k' 동안 반복되는 루프(while) 동안  'arr[lt] == 0' 여부를 검사하고 0이면 'chg--' 되어야 한다. 루프문이 실행될 때마다 늘 'lt++;' 되어야 한다.
3. lt가 변경되어야 할 때는 그 전에 누적된 cnt가 줄어들어야 한다.

 

개선 코드

import java.util.*;

public class Main {
    private int solution(int n, int k, int[] arr) {
        int answer = 0, lt = 0, chg = 0;
        for (int rt = 0; rt < n; rt++) {
            if (arr[rt] == 0) chg++;
            while (chg > k) {
                if (arr[lt++] == 0) chg--;
            }

            answer = Math.max(answer, rt - lt + 1);
        }

        return answer;
    }

    public static void main(String[] args) {
        Main al = new Main();
        Scanner sc = new Scanner(System.in);

        int n = sc.nextInt();
        int k = sc.nextInt();
        int[] arr = new int[n];
        for (int i = 0; i < n; i++) arr[i] = sc.nextInt();

        System.out.println(al.solution(n, k, arr));
    }
}

 0일 때 변경된 횟수를 올리고, 변경된 횟수가 주어진 횟수(k)보다 커지게 되면 같아질 때까지 lt를 당겨오는 방식이다. 물론 이 때 lt를 당겨오면서 옮기기 전의 lt가 0이었을 경우에는 0을 1로 변경한 횟수를 1 감소 시킨다. 초기에 작성한 코드보다 훨씬 간결해지고 불필요한 if문이 많이 제거된 것을 볼 수 있다.