package leetcode;

import java.util.ArrayList;
import java.util.List;

/**
 * https://leetcode.com/problems/max-consecutive-ones-iii/
 */
public class Problem1004 {
    public int longestOnes(int[] A, int K) {
        List<NumberCount> list = new ArrayList<>();
        for (int i = 0; i < A.length; i++) {
            if (A[i] == 0) {
                if (list.isEmpty()) {
                    list.add(new NumberCount(0, 1));
                    continue;
                }
                NumberCount nc = list.get(list.size() - 1);
                if (nc.number == A[i]) {
                    nc.count++;
                } else {
                    list.add(new NumberCount(0, 1));
                }
            } else {
                if (list.isEmpty()) {
                    list.add(new NumberCount(1, 1));
                    continue;
                }
                NumberCount nc = list.get(list.size() - 1);
                if (nc.number == A[i]) {
                    nc.count++;
                } else {
                    list.add(new NumberCount(1, 1));
                }
            }
        }
        int answer = 0;
        for (int i = 0; i < list.size(); i++) {
            int max = 0;
            int count = K;
            int j = i;
            while (j < list.size()) {
                NumberCount nc = list.get(j);
                if (nc.number == 0) {
                    if (count - nc.count >= 0) {
                        count -= nc.count;
                        max += nc.count;
                    } else {
                        max += count;
                        break;
                    }
                } else {
                    max += nc.count;
                }
                j++;
            }
            answer = Math.max(answer, max);
        }
        return answer;
    }

    private static class NumberCount {
        private final int number;
        private int count;

        public NumberCount(int number, int count) {
            this.number = number;
            this.count = count;
        }

        @Override
        public String toString() {
            return number + ": " + count;
        }
    }

    public static void main(String[] args) {
        Problem1004 prob = new Problem1004();
        System.out.println(prob.longestOnes(new int[]{1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0}, 2)); // 6
//        System.out.println(prob.longestOnes(new int[]{1, 1, 1}, 2)); // 3
//        System.out.println(prob.longestOnes(new int[]{0, 0, 0}, 3)); // 3
//        System.out.println(prob.longestOnes(new int[]{1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0}, 3)); // 10
//        System.out.println(prob.longestOnes(new int[]{0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1}, 3)); // 10
    }
}

@@ -1174,32 +1174,73 @@
 .size(); i++) %7B%0A
+            NumberCount n = list.get(i);%0A
             int 
@@ -1275,16 +1275,339 @@
 nt = K;%0A
+            if (n.number == 1) %7B%0A                max = n.count;%0A            %7D else %7B%0A                if (n.count - count %3E= 0) %7B%0A                    max = count;%0A                    count = 0;%0A                %7D else %7B%0A                    count -= n.count;%0A                    max = n.count;%0A                %7D%0A            %7D%0A
         
@@ -1619,16 +1619,20 @@
 nt j = i
+ + 1
 ;%0A      
@@ -2628,24 +2628,26 @@
 blem1004();%0A
+//
         Syst
@@ -3110,16 +3110,163 @@
 ; // 10%0A
+        System.out.println(prob.longestOnes(new int%5B%5D%7B1,0,0,0,1,1,0,0,1,1,0,0,0,0,0,0,1,1,1,1,0,1,0,1,1,1,1,1,1,0,1,0,1,0,0,1,1,0,1,1%7D, 8)); // 25%0A
     %7D%0A%7D%0A

Update problem 1004