Untitled

#include <iostream>
#include <queue>
#include <vector>
#include <unordered_map>
using namespace std;
 
const int MAXN = 20005; 
 
struct Node{
    int x;
    int y;
    int category;
    int distance;
    bool isDeleted;
} Node[MAXN];
 
int cnt;
int K, L;
unordered_map<int, int> ump; //  <mID:cus_Id>
vector<int> mat[41][41];
 
int dx[9] = { 0, -1, 0, 1, 0, -1, 1, -1, 1 };
int dy[9] = { -1, 0, 1, 0, 0, -1, 1, 1, -1 };
 
struct cmp{
    bool operator() (const int& a, const int& b) const {
        if (Node[a].distance == Node[b].distance) 
        {
            if (Node[a].x == Node[b].x)
                return Node[a].y > Node[b].y; 
            else
                return Node[a].x > Node[b].x; 
        }
        else
        {
            return Node[a].distance > Node[b].distance;
        }
    }
};
 
int getDistance(int x1, int y1, int x2, int y2) 
{ 
    return (abs(x1 - x2) + abs(y1 - y2)); 
}
 
bool isValid(int x, int y) 
{ 
    return (x >= 0 && y >= 0 && x < 41 && y < 41); 
}
 
void init(int _K, int _L)
{
    cnt = 0;
    K = _K;
    L = _L;
    
    ump.clear();
 
    for (int i = 0; i < 41; i++)
        for (int j = 0; j < 41; j++)
            mat[i][j].clear();
}
 
void addSample(int mID, int mX, int mY, int mC) // 20000
{
    ump[mID] = ++cnt;
 
    Node[cnt].x = mX;
    Node[cnt].y = mY;
    Node[cnt].category = mC;
    Node[cnt].distance = 0;
    Node[cnt].isDeleted = false;
 
    mat[mX / L][mY / L].push_back(cnt);
}
 
void deleteSample(int mID) // 1000
{
    Node[ump[mID]].isDeleted = true;
    ump.erase(mID);
}
 
int predict(int mX, int mY) // 10000
{
    priority_queue<int, vector<int>, cmp> pq; 
 
    int x = mX / L;
    int y = mY / L;
 
    for (int i = 0; i < 9; i++)
    {
        int X = x + dx[i];
        int Y = y + dy[i];
        if (isValid(X, Y))
        {
            for (auto itr : mat[X][Y])
            {
                if (!Node[itr].isDeleted)
                {
                    int currDist = getDistance(Node[itr].x, Node[itr].y, mX, mY);
                    Node[itr].distance = currDist;
                    if (currDist <= L) 
                        pq.push(itr);
                }
            }
        }
    }
 
    if (pq.size() < K) 
        return -1;
 
    int categoryCount[11] = { 0 };
 
    int sz = K;
    while (sz--) 
    {
        auto itr = pq.top();
        pq.pop();
 
        int ctgry = Node[itr].category;
        categoryCount[ctgry]++;
    }
 
    int ans = 0;
 
    for (int i = 1; i < 11; i++)
    {
        if (categoryCount[i] > categoryCount[ans])
            ans = i; 
    }
 
    return ans;
}