Untitled

import java.util.ArrayList;
import java.util.List;
import java.util.TreeSet;

class UserSolution {
    TreeSet<Passenger> sections[];
    Passenger passengers[];
    List<Integer> jobs[];

    void init(int N, int M, int J, int mPoint[], int mJobID[]) {
        sections = new TreeSet[N / M];
        jobs = new ArrayList[J]; // lưu id của các hành khách cùng job
        for (int i = 0; i < J; i++)
            jobs[i] = new ArrayList<>();
        passengers = new Passenger[N];
        for (int i = 0; i < N; i++) {
            if (sections[i / M] == null)
                sections[i / M] = new TreeSet<>(
                        (a, b) -> a.point == b.point ? a.id - b.id : b.point
                                - a.point); // sắp xếp giảm dần theo điểm và
                                            // tăng dần theo id
            passengers[i] = new Passenger(i, mPoint[i], mJobID[i], i / M);
            sections[i / M].add(passengers[i]);
            jobs[mJobID[i]].add(i);
        }
    }

    void destroy() {
    }

    int update(int mID, int mPoint) {
        Passenger p = passengers[mID];
        sections[p.sectionIndex].remove(p);
        p.point += mPoint;
        sections[p.sectionIndex].add(p);
        return p.point;
    }

    int updateByJob(int mJobID, int mPoint) {
        int sum = 0;
        for (int pID : jobs[mJobID]) {
            Passenger p = passengers[pID];
            sections[p.sectionIndex].remove(p);
            p.point += mPoint;
            sum += p.point;
            sections[p.sectionIndex].add(p);
        }
        return sum;
    }

    // section có kiểu dữ liệu treeSey đã đã được định nghĩa với comparator sắp
    // xếp
    // nên khi add dữ liệu sẽ tự động sắp xếp

    int move(int mNum) {
        int sum = 0;
        // lưu danh sách hành khách
        List<Passenger> temp[] = new ArrayList[sections.length - 1];
        for (int i = 0; i <= sections.length - 2; i++) {
            temp[i] = new ArrayList<>();
            // find and remove min from section 0 to length-2
            // store in temp[]
            for (int j = 0; j < mNum; j++) {
                Passenger smallest = sections[i].pollLast(); // do sắp xếp giảm
                                                             // dần
                smallest.sectionIndex++;
                temp[i].add(smallest);
                sum += smallest.point;
            }
            // find and remove max from section 1 to length-1
            // readd max to sections before
            for (int j = 0; j < mNum; j++) {
                Passenger largest = sections[i + 1].pollFirst();
                largest.sectionIndex--;
                sections[i].add(largest);
                sum += largest.point;
            }
        }
        // readd min to sections after
        for (int i = 0; i < temp.length; i++)
            sections[i + 1].addAll(temp[i]);
        return sum;
    }

    class Passenger {
        int id, point, jobID, sectionIndex;

        Passenger(int mID, int mPoint, int mJobID, int mSectionIndex) {
            id = mID;
            point = mPoint;
            jobID = mJobID;
            sectionIndex = mSectionIndex;
        }
    }
}