E. Yet Another Division Into Teams

Problem : 

There are n students at your university. The programming skill of the i-th student is ai. As a coach, you want to divide them into teams to prepare them for the upcoming ICPC finals. Just imagine how good this university is if it has 2⋅105 students ready for the finals!

Each team should consist of at least three students. Each student should belong to exactly one team. The diversity of a team is the difference between the maximum programming skill of some student that belongs to this team and the minimum programming skill of some student that belongs to this team (in other words, if the team consists of k students with programming skills a[i1],a[i2],…,a[ik], then the diversity of this team is maxj=1ka[ij]−minj=1ka[ij]).

The total diversity is the sum of diversities of all teams formed.

Your task is to minimize the total diversity of the division of students and find the optimal way to divide the students.

Input
The first line of the input contains one integer n (3≤n≤2⋅105) — the number of students.

The second line of the input contains n integers a1,a2,…,an (1≤ai≤109), where ai is the programming skill of the i-th student.

Output
In the first line print two integers res and k — the minimum total diversity of the division of students and the number of teams in your division, correspondingly.

In the second line print n integers t1,t2,…,tn (1≤ti≤k), where ti is the number of team to which the i-th student belong.

If there are multiple answers, you can print any. Note that you don't need to minimize the number of teams. Each team should consist of at least three students.


Topics : 
Sorting + DP + Greedy

Solution in Java : 


import java.util.*;
import java.lang.*;
import java.io.*;


public class Solution {

    static class Student implements Comparable<Student> {
        int programmingSkills;
        int index;

        Student(int programmingSkills, int index) {
            this.programmingSkills = programmingSkills;
            this.index = index;
        }

        public int compareTo(Student s) {
            return Integer.compare(programmingSkills, s.programmingSkills);
        }
    }

    static int INF = (int) 1e9 + 7;

    public static void main(String[] args) throws java.lang.Exception {
        out = new PrintWriter(new BufferedOutputStream(System.out));
        sc = new FastReader();

        int test = 1;
        for (int t = 1; t <= test; t++) {
            solve();
        }
        out.close();
    }

    private static void solve() {
        int n = sc.nextInt();
        Student[] students = new Student[n];

        for (int i = 0; i < n; i++) {
            int programmingSkills = sc.nextInt();
            students[i] = new Student(programmingSkills, i);
        }

        Arrays.sort(students);

        int[] teamSizes = new int[n + 1]; // size of the team in which ith student is in.
        Arrays.fill(teamSizes, -1);

        // dp[i] is the minimum total diversity we can get by dividing first i students into teams.
        int[] dp = new int[n + 1];
        Arrays.fill(dp, INF);
        dp[0] = 0;

        for (int i = 0; i < n; i++) { // for each student
            for (int j = 3; j <= 5 && j <= i + 1; j++) { // for each team size <= 5, since team of sizes > 5 can be divided into smaller teamSizes.
                if (dp[i + 1] > dp[i - j + 1] + students[i].programmingSkills - students[i - j + 1].programmingSkills) {
                    dp[i + 1] = dp[i - j + 1] + students[i].programmingSkills - students[i - j + 1].programmingSkills;
                    teamSizes[i + 1] = j;
                }
            }
        }

        int totalGroups = 0;
        int[] groups = new int[n];

        for (int i = n; i > 0; ) {
            totalGroups++;

            for (int j = i - teamSizes[i]; j < i; j++) {
                groups[students[j].index] = totalGroups;
            }

            i -= teamSizes[i];
        }

        out.println(dp[n] + " " + totalGroups);

        for (int i = 0; i < n; i++) {
            out.print(groups[i] + " ");
        }
        out.println();
    }


    public static FastReader sc;
    public static PrintWriter out;
    static class FastReader
    {
        BufferedReader br;
        StringTokenizer str;

        public FastReader()
        {
            br = new BufferedReader(new
                    InputStreamReader(System.in));
        }

        String next()
        {
            while (str == null || !str.hasMoreElements())
            {
                try
                {
                    str = new StringTokenizer(br.readLine());
                }
                catch (IOException  lastMonthOfVacation)
                {
                    lastMonthOfVacation.printStackTrace();
                }
            }
            return str.nextToken();
        }

        int nextInt()
        {
            return Integer.parseInt(next());
        }

        long nextLong()
        {
            return Long.parseLong(next());
        }

        double nextDouble()
        {
            return Double.parseDouble(next());
        }

        String nextLine()
        {
            String str = "";
            try
            {
                str = br.readLine();
            }
            catch (IOException lastMonthOfVacation)
            {
                lastMonthOfVacation.printStackTrace();
            }
            return str;
        }
    }

}