2611. 老鼠和奶酪

有两只老鼠和 n 块不同类型的奶酪，每块奶酪都只能被其中一只老鼠吃掉。

下标为 i 处的奶酪被吃掉的得分为：

如果第一只老鼠吃掉，则得分为 reward1[i] 。
如果第二只老鼠吃掉，则得分为 reward2[i] 。
给你一个正整数数组 reward1 ，一个正整数数组 reward2 ，和一个非负整数 k 。

请你返回第一只老鼠恰好吃掉 k 块奶酪的情况下，最大 得分为多少。

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/mice-and-cheese
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

class Solution {


    private void swap(Data[] datas, int a, int b) {
        Data t = datas[a];
        datas[a] = datas[b];
        datas[b] = t;
    }

    private int[] partition(Data[] datas, int L, int R) {
        int pivot = datas[(int) (L + (R - L) * Math.random())].sub;
        int left = L - 1, right = R + 1;
        int index = L;
        while (index < right) {
            if (datas[index].sub > pivot) {
                swap(datas, index++, ++left);
            } else if (datas[index].sub == pivot) {
                index++;
            } else {
                swap(datas, index, --right);
            }
        }
        return new int[]{left + 1, right - 1};
    }

    private int topK(Data[] datas, int k) {
        int left = 0, right = datas.length - 1;
        while (left <= right) {
            int[] partition = partition(datas, left, right);
            if (k >= partition[0] && k <= partition[1]) {
                return datas[k].sub;
            } else if (k < partition[0]) {
                right = partition[0] - 1;
            } else {
                left = partition[1] + 1;
            }
        }
        return datas[left].sub;
    }

    public int miceAndCheese(int[] reward1, int[] reward2, int k) {
        if (reward1 == null || reward1.length == 0 || reward2 == null || reward2.length == 0 || reward1.length != reward2.length) {
            return 0;
        }
        int n = reward1.length;
        int sum = 0;
        Data[] datas = new Data[n];
        for (int i = 0; i < n; i++) {
            sum += reward2[i];
            datas[i] = new Data(reward1[i] - reward2[i], reward1[i], reward2[i]);
        }
        topK(datas, k - 1);
        for (int i = 0; i < k; i++) {
            sum += datas[i].sub;
        }
        return sum;
    }
}

class Data {
    int sub;
    int reward1;
    int reward2;

    public Data(int sub, int reward1, int reward2) {
        this.sub = sub;
        this.reward1 = reward1;
        this.reward2 = reward2;
    }
}