Untitled

#include <algorithm>
#include <queue>
using namespace std;
#define MAX_N 100
#define MAX 101

int arr[MAX][MAX];
int bacteria[MAX][MAX];
int n;
int bacteria_type[3];
int visited[MAX][MAX] = { 0 };
int dx[4] = { -1, 0, 1, 0 };
int dy[4] = { 0, 1, 0, -1 };
int counter = 0;

struct Bacteria {
    int energy;
    int row;
    int col;
};

struct cmp {
    bool operator()(Bacteria& a, Bacteria& b) {
        if (a.energy == b.energy) {
            if (a.row == b.row) return a.col > b.col;
            return a.row > b.row;
        }
        return a.energy < b.energy;
    }
};

priority_queue<Bacteria, vector<Bacteria>, cmp> pq;

bool isValid(int x, int y) {
    return x > 0 && x <= n && y > 0 && y <= n;
}

void addDirection(int x, int y, int target) {
    queue<pair<int, int>> q;
    q.push({x, y});
    visited[x][y] = counter;

    while (!q.empty()) {
        auto [cur_x, cur_y] = q.front();
        q.pop();

        for (int i = 0; i < 4; i++) {
            int r_x = cur_x + dx[i];
            int c_y = cur_y + dy[i];

            if (!isValid(r_x, c_y) || visited[r_x][c_y] == counter) continue;

            if (bacteria[r_x][c_y] == 0) {
                pq.push({arr[r_x][c_y], r_x, c_y});
                visited[r_x][c_y] = counter;
            } else if (bacteria[r_x][c_y] == target) {
                visited[r_x][c_y] = counter;
                q.push({r_x, c_y});
            }
        }
    }
}

void init(int N, int mDish[MAX_N][MAX_N]) {
    n = N;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            arr[i + 1][j + 1] = mDish[i][j];
            bacteria[i + 1][j + 1] = 0;
        }
    }
    for (int i = 0; i < 3; i++) {
        bacteria_type[i] = 0;
    }
}

int dropMedicine(int mTarget, int mRow, int mCol, int mEnergy) {
    if (bacteria[mRow][mCol] != 0 && bacteria[mRow][mCol] != mTarget) {
        return bacteria_type[mTarget];
    }
    
    pq = priority_queue<Bacteria, vector<Bacteria>, cmp>();  // Reset priority queue
    counter++;

    if (bacteria[mRow][mCol] == 0) {
        mEnergy -= arr[mRow][mCol];
        bacteria[mRow][mCol] = mTarget;
        bacteria_type[mTarget]++;
    }

    pq.push({arr[mRow][mCol], mRow, mCol});
    visited[mRow][mCol] = counter;

    while (!pq.empty()) {
        auto c = pq.top();
        pq.pop();

        if (bacteria[c.row][c.col] == 0) {
            mEnergy -= c.energy;
            bacteria[c.row][c.col] = mTarget;
            bacteria_type[mTarget]++;
        }

        if (mEnergy <= 0) {
            return bacteria_type[mTarget];
        }

        if (bacteria[c.row][c.col] == mTarget) {
            addDirection(c.row, c.col, mTarget);
        }
    }

    return bacteria_type[mTarget];
}

int cleanBacteria(int mRow, int mCol) {
    int t = bacteria[mRow][mCol];
    if (t == 0) return -1;

    pq = priority_queue<Bacteria, vector<Bacteria>, cmp>();  // Reset priority queue
    counter++;

    pq.push({arr[mRow][mCol], mRow, mCol});
    visited[mRow][mCol] = counter;
    bacteria_type[t]--;

    while (!pq.empty()) {
        auto temp = pq.top();
        pq.pop();

        bacteria[temp.row][temp.col] = 0;

        for (int i = 0; i < 4; i++) {
            int x = temp.row + dx[i];
            int y = temp.col + dy[i];

            if (!isValid(x, y)) continue;

            if (visited[x][y] != counter && bacteria[x][y] == t) {
                pq.push({arr[x][y], x, y});
                visited[x][y] = counter;
                bacteria_type[t]--;
            }
        }
    }

    return bacteria_type[t];
}