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 C (GCC 5.4.1) (judge verdict: Wrong Answer)

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h> // uint64_t

#define max(a,b) ((a) > (b) ? (a) : (b))
#define min(a,b) ((a) > (b) ? (b) : (a))

int get_int(void) {
  int num;
  scanf("%d", &num);
  return num;
}

int get_int2(int *a1, int *a2) {
  scanf("%d %d", a1, a2);
  return 0;
}

struct sushi {
    int neta;
    int point;
};

#define NUM_MAX 100000
#define TYPES_MAX (NUM_MAX)

int desc_by_point(const void *a1, const void *a2) {
    struct sushi *s1 = (struct sushi*)a1;
    struct sushi *s2 = (struct sushi*)a2;
    return s2->point - s1->point;
}
#define HEAP_MAX (NUM_MAX+100)
#define ROOT_NODE 1
static struct sushi* heap[HEAP_MAX];
static int hidx = ROOT_NODE;

int comp(struct sushi *a1, struct sushi *a2) {
    return a1->point > a2->point;
}

void swap(struct sushi **a1, struct sushi **a2) {
    struct sushi *tmp = *a1;
    *a1 = *a2;
    *a2 = tmp;
}

// max heap
void enqueue(struct sushi *s) {
    int node = hidx;
    heap[hidx++] = s;
    int parent;
    while((parent = node/2)) {
        if(comp(heap[parent], heap[node])) break;
        swap(&heap[parent], &heap[node]);
        node = parent;
    }
    return;
}

struct sushi *delete_max(void) {
    struct sushi *ans = heap[ROOT_NODE];
    heap[ROOT_NODE] = heap[--hidx];
    int node = ROOT_NODE;
    while(1) {
        int largest = node;
        int left = node * 2;
        int right = node * 2 + 1;
        if(left < hidx && !comp(heap[node], heap[left])) {
            largest = left;
        }
        if(right < hidx && !comp(heap[node], heap[right])) {
            largest = right;
        }
        if(node == largest) break;
        swap(&heap[node], &heap[largest]);
        node = largest;
    }
    return ans;
}


static struct sushi* stack[NUM_MAX+50];
static int sidx = 0;

int is_empty(void) { return !sidx; }
void push(struct sushi *s) { stack[sidx++] = s; }
struct sushi *pop(void) { return stack[--sidx]; }


int main(void) {
    int num, k;
    get_int2(&num, &k);

    int i;
    static struct sushi sushis[NUM_MAX];
    for(i = 0; i < num; i++) {
        get_int2(&sushis[i].neta, &sushis[i].point);
    }
    qsort(sushis, num, sizeof(struct sushi), desc_by_point);

    static char checked[TYPES_MAX+1];
    int64_t base = 0;
    for(i = 0; i < num; i++) {
        struct sushi s = sushis[i];
        if(sidx >= k || checked[s.neta]) {
            enqueue(&sushis[i]);
            continue;
        }
        checked[s.neta]++;
        base += s.point;
        push(&sushis[i]);
    }
    // if (# of types) < K
    for(i = sidx; i < k; i++) {
        struct sushi *s = delete_max();
        base += s->point;
    }

#ifdef DEBUG
        printf("[stack: %d heap: %d] %lld\n", sidx, hidx, base);
#endif
    int64_t ans = base + sidx*sidx; // when x=max

    while(!is_empty()) {
        struct sushi *trash = pop();
        enqueue(trash);
        struct sushi *new = delete_max();
        base = base - trash->point + new->point;
#ifdef DEBUG
        printf("[stack: %d heap: %d] %lld\n", sidx, hidx, base);
#endif
        ans = max(ans, base + (int64_t)sidx*sidx);
    }

    printf("%lld\n", ans);
    return 0;
}