Problem:
Score : 400 points

Problem Statement

There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i.
You are choosing K among these N pieces to eat.
Your "satisfaction" here will be calculated as follows:

The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".

The base total deliciousness is the sum of the deliciousness of the pieces you eat.

The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.

You want to have as much satisfaction as possible.
Find this maximum satisfaction.

Constraints

1 \leq K \leq N \leq 10^5

1 \leq t_i \leq N

1 \leq d_i \leq 10^9

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

Output

Print the maximum satisfaction that you can obtain.

Sample Input 1

5 3
1 9
1 7
2 6
2 5
3 1

Sample Output 1

26

If you eat Sushi 1,2 and 3:

The base total deliciousness is 9+7+6=22.

The variety bonus is 2*2=4.

Thus, your satisfaction will be 26, which is optimal.

Sample Input 2

7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5

Sample Output 2

25

It is optimal to eat Sushi 1,2,3 and 4.

Sample Input 3

6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000

Sample Output 3

4900000016

Note that the output may not fit into a 32-bit integer type.

Solution in Java8 (OpenJDK 1.8.0) (judge verdict: Wrong Answer)



import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.List;
import java.util.stream.IntStream;


public class Main {
    public static void main(String[] args) {
        Scanner scanner=new Scanner();
        int n=scanner.nextInt();
        int k=scanner.nextInt();
        int[] t=new int[n];
        int[] d=new int[n];
        for(int i=0;i<n;i++){
            t[i]=scanner.nextInt();
            d[i]=scanner.nextInt();
        }
        List<Sushi> sushiList=new ArrayList<>();
        for(int i=0;i<n;i++){
            sushiList.add(new Sushi(t[i],d[i]));
        }
        //ネタでソート後おいしさでソート
        //後々PriorityQueueで(d,{1 (=同ネタ内で最高のおいしさ)or 0(=同ネタ内で別のスシのほうが最高のおいしさ)})というものを使うため
        Collections.sort(sushiList,Comparator.comparing((Sushi sushi)->sushi.t).thenComparing(sushi->sushi.d));
        //put(sushiList);
        Map<Sushi,FixedPair> map=new HashMap<>();
        for(int i=0;i<n-1;i++){
            if(sushiList.get(i).t!=sushiList.get(i+1).t){
                map.put(sushiList.get(i),new FixedPair(sushiList.get(i).d,1));
            }else{
                map.put(sushiList.get(i),new FixedPair(sushiList.get(i).d,0));
            }
        }
        map.put(sushiList.get(n-1),new FixedPair(sushiList.get(n-1).d,1));
        for(Map.Entry<Sushi,FixedPair> e:map.entrySet()){
            //put(format("%s->%s",e.getKey(),e.getValue()));
        }
        //おいしさ順にソート
        Collections.sort(sushiList,Comparator.comparingInt(s->s.d));
        //おいしさ降順にK個取った時のネタの種類数を数えるために使う
        Set<Integer> syurui0=new HashSet<>();
        //全てのネタの種類数を数えるために使う
        Set<Integer> syurui=new HashSet<>();
        long oishisaKiso=0;
        //既に使ったものを入れる
        //(?,0)で?の小さいものから使う
        PriorityQueue<FixedPair> used=new PriorityQueue<>(Comparator.comparing((FixedPair fp)->fp.getY()).thenComparing(fp->fp.getX()));
        //まだ使っていないものを入れる
        //(?,1)で?の大きいものから使う
        PriorityQueue<FixedPair> nonUsed=new PriorityQueue<>(Comparator.comparing((FixedPair fp)->-fp.getY()).thenComparing(fp->-fp.getX()));
        for(int i=n-1;i>=0;i--){
            syurui.add(sushiList.get(i).t);
            if(i>=n-k){
                syurui0.add(sushiList.get(i).t);
                oishisaKiso+= sushiList.get(i).d;
                used.add(map.get(sushiList.get(i)));
            }else{
                nonUsed.add(map.get(sushiList.get(i)));
            }
        }
        long[] f=new long[syurui.size()];
        for(int i=0;i<f.length;i++){
            f[i]=Long.MIN_VALUE;
        }
        f[syurui0.size()-1]=oishisaKiso;
        for(int i=syurui0.size();i<syurui.size();i++){
            //put(format("used:%s",used));
            //put(format("nonUsed:%s",nonUsed));
            FixedPair removeFP=used.poll();
            FixedPair nextFP=nonUsed.poll();
            if(nextFP==null) break;
            used.add(nextFP);
            f[i]=f[i-1]-removeFP.getX()+nextFP.getX();
        }
        //put(Arrays.toString(f));
        long max=Long.MIN_VALUE;
        for(int i=0;i<f.length;i++){
            max=max(max,f[i]+(long)Math.pow(i+1,2));
        }
        put(max);



    }
    static class Sushi{
        final int t;
        final int d;
        public Sushi(int t,int d){
            this.t=t;
            this.d=d;
        }

        @Override
        public String toString() {
            return "Sushi{" +
                    "t=" + t +
                    ", d=" + d +
                    '}';
        }
    }





    public static int upper_bound(int[] a,int val){
        return upper_bound(a,0,a.length,val);
    }
    public static int upper_bound(int[] a,int l,int r,int val){
        if(r-l==1){
            if(a[l]>val) return l;
            return r;
        }
        int mid=(l+r)/2;
        if(a[mid]>val){
            return upper_bound(a,l,mid,val);
        }else{
            return upper_bound(a,mid,r,val);
        }
    }
    public static <T> int upper_bound(T[] a,int l,int r,T val,Comparator<T> comparator){
        if(r-l==1){
            if(comparator.compare(a[l],val)>0) return l;
            return r;
        }
        int mid=(l+r)/2;
        if(comparator.compare(a[mid],val)>0){
            return upper_bound(a,l,mid,val,comparator);
        }else{
            return upper_bound(a,mid,r,val,comparator);
        }
    }
    public static int lower_bound(int[] a,int val){
        return lower_bound(a,0,a.length,val);
    }
    public static int lower_bound(int[] a,int l,int r,int val){
        if(r-l==1){
            if(a[l]<val) return r;
            return l;
        }
        int mid=(l+r)/2;
        if(a[mid]<val){
            return lower_bound(a,mid,r,val);
        }else{
            return lower_bound(a,l,mid,val);
        }

    }

    public static <T> int lower_bound(T[] a,int l,int r,T val,Comparator<T> comparator){
        if(r-l==1){
            if(comparator.compare(a[l],val)<0) return r;
            return l;
        }
        int mid=(l+r)/2;
        if(comparator.compare(a[mid],val)<0){
            return lower_bound(a,mid,r,val,comparator);
        }else{
            return lower_bound(a,l,mid,val,comparator);
        }
    }


    public static int max(int a,int b){
        return Math.max(a,b);
    }
    public static long max(long a,long b){
        return Math.max(a,b);
    }
    public static int min(int a,int b){
        return Math.min(a,b);
    }
    public static long min(long a,long b){
        return Math.min(a,b);
    }
    public static int abs(int a){
        return Math.abs(a);
    }
    public static long abs(long a){
        return Math.abs(a);
    }
    public static void print(Object object){
        System.out.print(object);
    }
    public static void put(Object object) {
        System.out.println(object);
    }
    public static void put(){
        System.out.println();
    }
    public static String format(String string, Object... args) {
        return String.format(string, args);
    }
}

final class Scanner {
    private final InputStream in = System.in;
    private final byte[] buffer = new byte[1024];
    private int ptr = 0;
    private int buflen = 0;

    private boolean hasNextByte() {
        if (ptr < buflen) {
            return true;
        } else {
            ptr = 0;
            try {
                buflen = in.read(buffer);
            } catch (IOException e) {
                e.printStackTrace();
            }
            if (buflen <= 0) {
                return false;
            }
        }
        return true;
    }

    private int readByte() {
        if (hasNextByte())
            return buffer[ptr++];
        else
            return -1;
    }

    private boolean isPrintableChar(int c) {
        return 33 <= c && c <= 126;
    }

    public boolean hasNext() {
        while (hasNextByte() && !isPrintableChar(buffer[ptr]))
            ptr++;
        return hasNextByte();
    }

    public String next() {
        if (!hasNext())
            throw new NoSuchElementException();
        StringBuilder sb = new StringBuilder();
        int b = readByte();
        while (isPrintableChar(b)) {
            sb.appendCodePoint(b);
            b = readByte();
        }
        return sb.toString();
    }

    public long nextLong() {
        if (!hasNext())
            throw new NoSuchElementException();
        long n = 0;
        boolean minus = false;
        int b = readByte();
        if (b == '-') {
            minus = true;
            b = readByte();
        }
        if (b < '0' || '9' < b) {
            throw new NumberFormatException();
        }
        while (true) {
            if ('0' <= b && b <= '9') {
                n *= 10;
                n += b - '0';
            } else if (b == -1 || !isPrintableChar(b)) {
                return minus ? -n : n;
            } else {
                throw new NumberFormatException();
            }
            b = readByte();
        }
    }

    public int nextInt() {
        long nl = nextLong();
        if (nl < Integer.MIN_VALUE || nl > Integer.MAX_VALUE)
            throw new NumberFormatException();
        return (int) nl;
    }

    public double nextDouble() {
        return Double.parseDouble(next());
    }
}
final class FixedPair {
    final private long x, y;
    final static public FixedPair ZEROS=new FixedPair(0,0);
    FixedPair(long x, long y) {
        this.x = x;
        this.y = y;
    }

    public long getX() {
        return x;
    }

    public long getY() {
        return y;
    }

    @Override
    public int hashCode() {
        return (int)x+(int)y;
    }

    @Override
    public boolean equals(Object obj) {
        boolean result=super.equals(obj);
        if(obj.getClass()!=this.getClass()){
            return false;
        }
        FixedPair pair=(FixedPair)obj;
        if(this.x==pair.x&&this.y==pair.y) return true;
        return false;
    }

    @Override
    public String toString() {
        return String.format("(%d,%d)", x, y);
    }
}
final class Tuple<T,V>{
    //immutabl1でないことに注意(T,Vがmutableの場合変更可能)
    final private T t;
    final private V v;
    Tuple(T t,V v){
        this.t=t;
        this.v=v;
    }

    public T getT() {
        return t;
    }

    public V getV() {
        return v;
    }

    @Override
    public int hashCode() {
        return (t.hashCode()+v.hashCode());
    }

    @Override
    @SuppressWarnings("unchecked")
    public boolean equals(Object obj) {
        if(obj.getClass()!=this.getClass()){
            return false;
        }
        Tuple<T,V> tuple=(Tuple<T, V>) obj;
        return tuple.t.equals(this.t)&&tuple.v.equals(this.v);
    }

    @Override
    public String toString() {
        return String.format("<Tuple>=<%s,%s>",t,v);
    }
}
final class LowerBoundComparator<T extends Comparable<? super T>>
        implements Comparator<T>
{
    public int compare(T x, T y)
    {
        return (x.compareTo(y) >= 0) ? 1 : -1;
    }
}

final class UpperBoundComparator<T extends Comparable<? super T>>
        implements Comparator<T>
{
    public int compare(T x, T y)
    {
        return (x.compareTo(y) > 0) ? 1 : -1;
    }
}
final class Util {
    static long gcd(long a,long b){
        //最大公約数　
        if(a%b==0)return b;
        return gcd(b,a%b);
    }
    static long lcm(long a,long b){
        //最小公倍数
        long gcd=gcd(a,b);
        long result=b/gcd;
        return a*result;
    }
    static long kaijoMod(int n,int mod){
        if(n<1) return -1;
        long result=1;
        for(int i=n;i>1;i--){
            result*=i;
            result%=mod;
        }
        return result;
    }
    static <T extends Comparable<T>> Map<T,Integer> count(List<T> list){
        //副作用
        Collections.sort(list);
        Map<T,Integer> result=new HashMap<>();
        int l=0,r=0;
        while(l<list.size()){
            while(r<list.size()-1&&list.get(r).equals(r+1)){
                r++;
            }
            result.put(list.get(r),r-l+1);
            r++;
            l=r;
        }
        return result;

    }
    static Map<Integer,Integer> count(int[] array){
        //副作用
        Arrays.sort(array);
        Map<Integer,Integer> result=new HashMap<>();
        int l=0,r=0;
        while(l<array.length){
            while(r<array.length-1&&array[r]==array[r+1]){
                r++;
            }
            result.put(array[l],r-l+1);
            r++;
            l=r;
        }
        return result;
    }
    static <T> String toStringBWS(Iterable<T> iterable){
        Iterator<T> ite=iterable.iterator();
        return toStringBWS(ite);
    }
    static <T> String toStringBWS(Iterator<T> ite){
        StringBuilder sb=new StringBuilder();
        sb.append(ite.next());
        while(ite.hasNext()){
            sb.append(" ");
            sb.append(ite.next());
        }
        return sb.toString();
    }
    static int[] factoringInPrimeNumbers(long n,int size){
        //素因数分解
        //sizeがnに比べて小さい場合完全に素因数分解出来ていない可能性がある
        int[] result=new int[size];
        for(int i=2;n>1&&i< result.length;i++){
            while(n%i==0){
                result[i]++;
                n/=i;
            }
        }
        return result;

    }
}