Untitled

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

class UserSolution2 {
    class Passenger {
        int id;
        int jId;
        int point;

        public Passenger(int id, int jId, int point) {
            this.id = id;
            this.jId = jId;
            this.point = point;
        }
    }

    ArrayList<PriorityQueue<Passenger>> arrQueue;
    HashMap<Integer, Passenger> mapP;
    HashMap<Integer, ArrayList<Integer>> jobType;

    int numP, numS, numJ, numPS;

    void init(int N, int M, int J, int mPoint[], int mJobID[]) {
        numP = N;
        numS = N / M;
        numPS = M;
        numJ = J;
        mapP = new HashMap<>();
        jobType = new HashMap<>();
        arrQueue = new ArrayList<>();

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

        for (int i = 0; i < numS; i++) {
            PriorityQueue<Passenger> PQ = new PriorityQueue<>((a, b) -> a.point == b.point ? b.id - a.id : a.point - b.point);
            for (int j = i * M; j < (i + 1) * M; j++) {
                PQ.add(mapP.get(j));
            }
            arrQueue.add(PQ);
        }
    }

    void destroy() {
    }

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

    int updateByJob(int mJobID, int mPoint) {
        int sum = 0;
        if (jobType.containsKey(mJobID)) {
            for (int i : jobType.get(mJobID)) {
                mapP.get(i).point += mPoint;
                sum += mapP.get(i).point;
            }
        }
        return sum;
    }

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

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

        return sum;
    }
}