Untitled

import java.util.ArrayList;
import java.util.HashMap;
import java.util.PriorityQueue;

class UserSolution {
    class Passenger {
        int pId;
        int jId;
        int pPoint;

        public Passenger(int pId, int jId, int pPoint) {
            this.pId = pId;
            this.jId = jId;
            this.pPoint = pPoint;
        }
    }

    ArrayList<PriorityQueue<Passenger>> arrMin;
    HashMap<Integer, Passenger> mapPass;
    HashMap<Integer, ArrayList<Integer>> jtype;

    int NOP;
    int NOS;
    int TOJ;
    int NOM;

    void init(int N, int M, int J, int mPoint[], int mJobID[]) {
        NOP = N;
        NOS = N / M;
        NOM = M;
        TOJ = J;
        mapPass = new HashMap<>();
        jtype = new HashMap<>();
        arrMin = new ArrayList<>();

        for (int i = 0; i < N; i++) {
            Passenger P = new Passenger(i, mJobID[i], mPoint[i]);
            mapPass.put(i, P);
            jtype.computeIfAbsent(mJobID[i], k -> new ArrayList<>()).add(i);
        }

        for (int i = 0; i < NOS; i++) {
            PriorityQueue<Passenger> PQ = new PriorityQueue<>((a, b) -> a.pPoint == b.pPoint ? b.pId - a.pId : a.pPoint - b.pPoint);
            for (int j = i * M; j < (i + 1) * M; j++) {
                PQ.add(mapPass.get(j));
            }
            arrMin.add(PQ);
        }
    }

    void destroy() {
    }

    int update(int mID, int mPoint) {
        if (mapPass.containsKey(mID)) {
            mapPass.get(mID).pPoint += mPoint;
        }
        return mapPass.get(mID).pPoint;
    }

    int updateByJob(int mJobID, int mPoint) {
        int sum = 0;
        if (jtype.containsKey(mJobID)) {
            for (int i : jtype.get(mJobID)) {
                mapPass.get(i).pPoint += mPoint;
                sum += mapPass.get(i).pPoint;
            }
        }
        return sum;
    }

    int move(int mNum) {
        int sum = 0;

        Passenger[][] Pmax = new Passenger[NOS - 1][mNum];
        Passenger[][] Pmin = new Passenger[NOS - 1][mNum];
        
        for (int i = 0; i < NOS - 1; i++) {
            arrMin.set(i, new PriorityQueue<>(arrMin.get(i)));
            for (int k = 0; k < mNum; k++) {
                Pmin[i][k] = arrMin.get(i).poll();
            }
        }
        
        for (int i = 1; i < NOS; i++) {
            arrMin.set(i, new PriorityQueue<>((a, b) -> a.pPoint == b.pPoint ? a.pId - b.pId : b.pPoint - a.pPoint));
            for (int k = 0; k < mNum; k++) {
                Pmax[i-1][k] = arrMin.get(i).poll();
            }
        }
        
        for (int i = 0; i < NOS - 1; i++) {
            for (int k = 0; k < mNum; k++) {
                arrMin.get(i).add(Pmax[i][k]);
                arrMin.get(i + 1).add(Pmin[i][k]);
                sum += Pmin[i][k].pPoint + Pmax[i][k].pPoint;
            }
        }

        return sum;
    }
}