Untitled

import java.util.*;
 
class UserSolution {
 
    HashMap<Integer, Passenger> passengers;
    HashMap<Integer, HashSet<Passenger>> jobs;
    TreeSet<Passenger>[] sections;
 
    void init(int N, int M, int J, int mPoint[], int mJobID[])
    {
        passengers = new HashMap<>();
        jobs = new HashMap<>();
 
        sections = new TreeSet[N/M];
        for (int i = 0; i < sections.length; i++) {
            sections[i] = new TreeSet<>(new Comparator<Passenger>() {
                @Override
                public int compare(Passenger o1, Passenger o2) {
                    if (o1.point == o2.point) {
                        return o1.id-o2.id;
                    }
                    return o2.point-o1.point;
                }
            });
        }
        for (int i = 0; i < N; i++) {
            Passenger new_passenger = new Passenger(i, i/M, mJobID[i], mPoint[i]);
            sections[i/M].add(new_passenger);
            passengers.put(i, new_passenger);
             
            HashSet<Passenger> set = jobs.computeIfAbsent(mJobID[i], k -> new HashSet<>());
            set.add(new_passenger);
        }
    }
 
    void destroy()
    {
    }
 
    int update(int mID, int mPoint)
    {
        Passenger edited = passengers.get(mID);
        sections[edited.section].remove(edited);
        edited.point += mPoint;
        sections[edited.section].add(edited);
        return edited.point;
    }
 
    int updateByJob(int mJobID, int mPoint)
    {
        HashSet<Passenger> group = jobs.get(mJobID);
        int sum = 0;
        for (Passenger passenger : group) {
            sections[passenger.section].remove(passenger);
            passenger.point += mPoint;
            sum += passenger.point;
            sections[passenger.section].add(passenger);
        }
        return sum;
    }
 
    int move(int mNum)
    {
        int sum = 0;
        ArrayList<Passenger>[] temp = new ArrayList[sections.length];
        for (int i = 0; i < temp.length; i++) {
            temp[i] = new ArrayList<>();
        }
        for (int i = 0; i < sections.length; i++) {
            // ArrayList<Passenger> front = new ArrayList<>();
            // ArrayList<Passenger> back = new ArrayList<>();
            for (int j = 0; j < mNum; j++) {
                if (i > 0) {
                    Passenger x = sections[i].pollFirst();
                    x.section --;
                    // front.add(x);
                    sum += x.point;
                    temp[i-1].add(x);
                }
                if (i < sections.length-1) {
                    Passenger y = sections[i].pollLast();
                    y.section ++;
                    // back.add(y);
                    sum += y.point;
                    temp[i+1].add(y);    
                } 
            }
        }
        for (int i = 0; i < sections.length; i++) {
            sections[i].addAll(temp[i]);
        }
        return sum;
    }
}
 
class Passenger {
    int id;
    int section;
    int job;
    int point;
 
    Passenger(int id, int section, int job, int point){
        this.id = id;
        this.section = section;
        this.job = job;
        this.point = point;
    }
 
}