package leetcode;

import java.util.Comparator;
import java.util.TreeSet;

/**
 * https://leetcode.com/problems/partition-array-for-maximum-sum/
 */
public class Problem1043 {
    public int maxSumAfterPartitioning(int[] A, int K) {
        TreeSet<Integer> set = new TreeSet<>(Comparator.reverseOrder());
        for (int a : A) {
            set.add(a);
            if (set.size() > K) {
                set.pollLast();
            }
        }
        int answer = 0;
        int prev = 0;
        for (int i = 0; i < A.length; i++) {
            if (set.contains(A[i])) {
                answer += (i - prev) * A[i];
                prev = i;
            }
        }
        return answer;
    }

    public static void main(String[] args) {
        Problem1043 prob = new Problem1043();
        System.out.println(prob.maxSumAfterPartitioning(new int[]{1, 15, 7, 9, 2, 5, 10}, 3)); // 84
        System.out.println(prob.maxSumAfterPartitioning(new int[]{1, 15, 7, 10, 2, 5, 9}, 3)); // 84
        System.out.println(prob.maxSumAfterPartitioning(new int[]{1, 15, 7, 10, 2, 5, 9}, 4)); // 87
        System.out.println(prob.maxSumAfterPartitioning(new int[]{10, 15, 7, 9, 2, 5, 1}, 4)); // 81
        System.out.println(prob.maxSumAfterPartitioning(new int[]{10, 15, 7, 9, 2, 5, 1}, 2)); // 68
    }
}

@@ -16,64 +16,8 @@
 e;%0A%0A
-import java.util.Comparator;%0Aimport java.util.TreeSet;%0A%0A
 /**%0A
@@ -182,71 +182,238 @@
     
-T
 re
-eSet%3CInteger%3E set = new TreeSet%3C%3E(Comparator.reverseOrder())
+turn maxSumAfterPartitioning(A, K, 0);%0A    %7D%0A%0A    private static int maxSumAfterPartitioning(int%5B%5D array, int k, int index) %7B%0A        if (index %3E= array.length) %7B%0A            return 0;%0A        %7D%0A        int answer = array%5Bindex%5D
 ;%0A  
@@ -431,13 +431,25 @@
 int 
-a : A
+i = 1; i %3C k; i++
 ) %7B%0A
@@ -464,18 +464,19 @@
     
-set.add(a)
+int val = 0
 ;%0A  
@@ -493,22 +493,33 @@
 if (
-set.size() %3E K
+index + i %3E= array.length
 ) %7B%0A
@@ -538,22 +538,13 @@
     
-set.pollLast()
+break
 ;%0A  
@@ -551,38 +551,32 @@
           %7D%0A    
-    %7D%0A
         int answ
@@ -571,22 +571,19 @@
     int 
-answer
+max
  = 0;%0A  
@@ -588,26 +588,8 @@
     
-    int prev = 0;%0A
     
@@ -605,30 +605,35 @@
 int 
-i = 0; i %3C A.length; i
+j = index; j %3C index + i; j
 ++) 
@@ -650,33 +650,42 @@
     
-if (set.contains(A%5Bi%5D)) %7B
+    max = Math.max(max, array%5Bj%5D);
 %0A   
@@ -697,80 +697,135 @@
     
+%7D%0A
     
-answer += (i - prev) * A%5Bi%5D;%0A                prev = i;%0A            %7D
+        val = maxSumAfterPartitioning(array, k, index + i) + (max * (i + 1));%0A            answer = Math.max(answer, val);
 %0A   
@@ -952,16 +952,200 @@
 1043();%0A
+//        System.out.println(prob.maxSumAfterPartitioning(new int%5B%5D%7B1, 2, 3, 4%7D, 1)); // 10%0A        System.out.println(prob.maxSumAfterPartitioning(new int%5B%5D%7B1, 2, 3, 4%7D, 3)); // 13%0A//
         
@@ -1229,32 +1229,34 @@
 10%7D, 3)); // 84%0A
+//
         System.o
@@ -1340,16 +1340,18 @@
 ; // 84%0A
+//
         
@@ -1443,16 +1443,18 @@
 ; // 87%0A
+//
         
@@ -1546,16 +1546,18 @@
 ; // 81%0A
+//
         

Update problem 1043