HZAUoj 1016: Array C 【贪心】

题目链接：HZAUoj 1016: Array C 【贪心】

#include <cstdio>
#include <cstring>
#include <cmath>
#include <map>
#include <set>
#include <iostream>
#include <algorithm>
#include <queue>
#include <string>
#include <stack>
#define CLR(a, b) memset(a, b, sizeof(a))
using namespace std;
typedef long long LL;
typedef pair<int, int> pii;
const int MAXN = 3 * 1e6 + 10;
const int MOD = 9973;
const int INF = 0x3f3f3f3f;
struct Node {
    int id, cost;
    friend bool operator < (Node a, Node b) {
        return a.cost < b.cost;
    }
};
int a[1010], b[1010];
int main()
{
    int n, m;
    while(scanf("%d%d", &n, &m) != EOF) {
        int sum = 0;
        for(int i = 1; i <= n; i++) {
            scanf("%d", &a[i]);
            sum += a[i];
        }
        int sub = sum - m;
        priority_queue<Node> Q; Node now;
        for(int i = 1; i <= n; i++) {
            scanf("%d", &b[i]);
            now.id = i; now.cost = (2 * a[i] - 1) * b[i];
            Q.push(now);
        }
        LL ans = 0;
        while(sub--) {
            now = Q.top(); Q.pop(); a[now.id]--;
            now.cost = (2 * a[now.id] - 1) * b[now.id];
            Q.push(now);
        }
        for(int i = 1; i <= n; i++) {
            ans += a[i] * a[i] * b[i];
        }
        printf("%lld\n", ans);
    }
    return 0;
}