Untitled

#define MAX_N 100
#include <queue>
#include <iostream>
#include <vector>
using namespace std;

vector<vector<int>> dish(MAX_N, vector<int>(MAX_N));  // energy values
vector<vector<int>> bacteriaType(MAX_N, vector<int>(MAX_N, 0));  // bacteria types
vector<vector<int>> visited(MAX_N, vector<int>(MAX_N, 0));
int n, point = 0;
int cnt[2] = {0, 0};  // count of bacteria types

const int dx[] = {0, 0, 1, -1};
const int dy[] = {1, -1, 0, 0};

class CompareNodes {
public:
    bool operator()(const pair<int, pair<int, int>>& a, const pair<int, pair<int, int>>& b) {
        int energy1 = a.first;
        int energy2 = b.first;
        int row1 = a.second.first;
        int row2 = b.second.first;
        int col1 = a.second.second;
        int col2 = b.second.second;
        
        if (energy1 != energy2) return energy1 < energy2;
        if (row1 != row2) return row1 > row2;
        return col1 > col2;
    }
};

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

void init(int N, int mDish[MAX_N][MAX_N]) {
    n = N;
    cnt[0] = cnt[1] = point = 0;
    
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            dish[i][j] = mDish[i][j];
            bacteriaType[i][j] = 0;
            visited[i][j] = 0;
        }
    }
}

int dropMedicine(int mTarget, int mRow, int mCol, int mEnergy) {
    int r = mRow - 1;
    int c = mCol - 1;
    
    if (bacteriaType[r][c] != 0 && bacteriaType[r][c] != mTarget) {
        return cnt[mTarget - 1];
    }

    if (bacteriaType[r][c] == 0) {
        bacteriaType[r][c] = mTarget;
        cnt[mTarget - 1]++;
        mEnergy -= dish[r][c];
    }
    
    point++;
    priority_queue<pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, CompareNodes> pq;
    queue<pair<int, int>> bfs;
    
    visited[r][c] = point;
    bfs.push({r, c});
    
    while (mEnergy > 0) {
        while (!bfs.empty()) {
            auto [currRow, currCol] = bfs.front();
            bfs.pop();
            
            for (int i = 0; i < 4; i++) {
                int newRow = currRow + dx[i];
                int newCol = currCol + dy[i];
                
                if (isValid(newRow, newCol) && visited[newRow][newCol] < point) {
                    visited[newRow][newCol] = point;
                    if (bacteriaType[newRow][newCol] == 0) {
                        pq.push({dish[newRow][newCol], {newRow, newCol}});
                    }
                    else if (bacteriaType[newRow][newCol] == mTarget) {
                        bfs.push({newRow, newCol});
                    }
                }
            }
        }
        
        if (pq.empty()) break;
        
        auto [energy, pos] = pq.top();
        pq.pop();
        auto [newRow, newCol] = pos;
        
        bacteriaType[newRow][newCol] = mTarget;
        cnt[mTarget - 1]++;
        mEnergy -= dish[newRow][newCol];
        bfs.push({newRow, newCol});
    }
    
    return cnt[mTarget - 1];
}

int cleanBacteria(int mRow, int mCol) {
    int r = mRow - 1;
    int c = mCol - 1;
    int type = bacteriaType[r][c];
    
    if (type == 0) return -1;
    
    point++;
    queue<pair<int, int>> bfs;
    visited[r][c] = point;
    cnt[type - 1]--;
    bacteriaType[r][c] = 0;
    bfs.push({r, c});
    
    while (!bfs.empty()) {
        auto [currRow, currCol] = bfs.front();
        bfs.pop();
        
        for (int i = 0; i < 4; i++) {
            int newRow = currRow + dx[i];
            int newCol = currCol + dy[i];
            
            if (isValid(newRow, newCol) && visited[newRow][newCol] < point && 
                bacteriaType[newRow][newCol] == type) {
                bacteriaType[newRow][newCol] = 0;
                cnt[type - 1]--;
                visited[newRow][newCol] = point;
                bfs.push({newRow, newCol});
            }
        }
    }
    
    return cnt[type - 1];
}