Untitled

#include<iostream>
#include<unordered_map>
#include<queue>
#define MAX_N    105
#define MAX_TOOL 5005
using namespace std;
 
class TOOL {
public:
    int toolidx;
    int stepno;
    int Pt;
    int finishT;
};
TOOL tools[MAX_TOOL];
class cmp {
public:
    bool operator()(int a, int b)const {
        return tools[a].Pt > tools[b].Pt;
    }
};
class cmp2 {
public:
    bool operator()(int x, int y)const {
        if (tools[x].finishT == tools[y].finishT) {
            if (tools[x].stepno == tools[y].stepno) {
                return tools[x].Pt > tools[y].Pt;
            }
            else {
                return tools[x].stepno < tools[y].stepno;
            }
        }
        return tools[x].finishT > tools[y].finishT;
    }
};
priority_queue <int, vector<int>, cmp> stepPq[MAX_N]; //stepwise queue
priority_queue <int, vector<int>, cmp2> timePq;  // all the tools in processing
int n;
 
// for each step
int total[MAX_N];
int waiting[MAX_N];
 
void init(int N) {
    n = N;
    for (int i = 0; i <= N; i++) {
        stepPq[i] = priority_queue<int, vector<int>, cmp>();
         
        //because step is empty at start
        total[i] = waiting[i] = 0;
    }
    timePq = priority_queue<int, vector<int>, cmp2>();
}
 
void setupTool(int T, int stepNo[5000], int procTime[5000]) {
    for (int i = 0; i < T; i++) {
        tools[i].toolidx = i;
        tools[i].stepno = stepNo[i];
        tools[i].Pt = procTime[i];
        tools[i].finishT = 0;
 
        // add tool to priority_queue of step i;
        stepPq[stepNo[i]].push(i);
    }
 
}
void assign_process(int time, int stepno) {
    while (stepPq[stepno].size() != 0 && waiting[stepno] != 0) //priority q has tools and lots are in waiting/standby
    {
        // get min time tool from PQ
        int tool_id = stepPq[stepno].top();
        stepPq[stepno].pop();
         
        // assign this tool
        tools[tool_id].finishT = time + tools[tool_id].Pt;
 
        // update waiting and process_pq
        timePq.push(tool_id);
        waiting[stepno]--;
    }
}
 
void update(int time)
{
    // for updating time
    // time pq has tools in process and finish time of top < current time 
    while (timePq.size() > 0 && tools[timePq.top()].finishT <= time) { 
        // remove from time pq 
        int tool_id = timePq.top();
        timePq.pop();
 
        // add to step pq
        int stepno = tools[tool_id].stepno;
        int current_time = tools[tool_id].finishT;
        total[stepno]--;
 
        if (waiting[stepno] > 0) // a lot is waiting in this step
        {
            // assign this tool to that lot
            // a process finished
             
 
            tools[tool_id].finishT = current_time + tools[tool_id].Pt;
 
            // update waiting and process_pq
            timePq.push(tool_id);
            waiting[stepno]--;
 
 
        }
        else // no waiting => simply add to step PQ
        {
            stepPq[stepno].push(tool_id);
        }
        //move to next step
        stepno++;
        total[stepno]++;
        //updating total and waiting
        if (stepno != n) {
            waiting[stepno]++;
            if (stepPq[stepno].size() > 0) {
                assign_process(current_time, stepno);
            }
        }
    }
 
}
 
void addLot(int time, int number) {
    update(time); // update before adding
    waiting[0] += number;
    total[0] += number;
    assign_process(time,0);
}
 
int simulate(int time, int wip[MAX_N]) {
    update(time); 
    // after update we have updated step priority queue and process pq
    for (int i = 0; i < n; i++) {
        wip[i] = total[i];
    }
 
    return total[n];
}