Untitled

#include <iostream>
#include <vector>
#include <unordered_map>
#include <set>

using namespace std;

set<int> rowMap[3001];
set<int> colMap[3001];
set<int> leftDiag[6001];
set<int> rightDiag[6001];

struct Hole {
    int row, col, cnt;
    bool check;
};

Hole holes[30001];

struct Diagonal {
    int sr, sc, er, ec;
};

Diagonal leftDiagonal[6001];
Diagonal rightDiagonal[6001];
int n;
int step = 0;

void init(int N) {
    n = N;
    for (int i = 0; i < 3001; i++) {
        rowMap[i].clear();
        colMap[i].clear();
    }
    for (int i = 0; i < 6001; i++) {
        leftDiag[i].clear();
        rightDiag[i].clear();
    }
    for (int i = 0; i < 30001; i++) {
        holes[i].row = -1;
        holes[i].col = -1;
        holes[i].check = false;
        holes[i].cnt = 0;
    }
    for (int i = 0; i < 6001; i++) {
        leftDiagonal[i].sc = -1;
        leftDiagonal[i].sr = -1;
        leftDiagonal[i].ec = -1;
        leftDiagonal[i].er = -1;

        rightDiagonal[i].sc = -1;
        rightDiagonal[i].sr = -1;
        rightDiagonal[i].ec = -1;
        rightDiagonal[i].er = -1;
    }
}

void addDiagonal(int mARow, int mACol, int mBRow, int mBCol) {
    if (mACol + mARow == mBCol + mBRow) { // Left diagonal
        if (mARow < mBRow) {
            leftDiagonal[mARow + mACol].sr = mARow;
            leftDiagonal[mARow + mACol].sc = mACol;
            leftDiagonal[mARow + mACol].er = mBRow;
            leftDiagonal[mARow + mACol].ec = mBCol;
        } else {
            leftDiagonal[mARow + mACol].sr = mBRow;
            leftDiagonal[mARow + mACol].sc = mBCol;
            leftDiagonal[mARow + mACol].er = mARow;
            leftDiagonal[mARow + mACol].ec = mACol;
        }
    } else {
        if (mARow < mBRow) {
            rightDiagonal[mARow - mACol + n].sr = mARow;
            rightDiagonal[mARow - mACol + n].sc = mACol;
            rightDiagonal[mARow - mACol + n].er = mBRow;
            rightDiagonal[mARow - mACol + n].ec = mBCol;
        } else {
            rightDiagonal[mARow - mACol + n].sr = mBRow;
            rightDiagonal[mARow - mACol + n].sc = mBCol;
            rightDiagonal[mARow - mACol + n].er = mARow;
            rightDiagonal[mARow - mACol + n].ec = mACol;
        }
    }
}

void addHole(int mRow, int mCol, int mID) {
    holes[mID].row = mRow;
    holes[mID].col = mCol;
    holes[mID].check = true;

    rowMap[mRow].insert(mID);
    colMap[mCol].insert(mID);
    leftDiag[mRow + mCol].insert(mID);
    rightDiag[mRow - mCol + n].insert(mID);
}

void eraseHole(int mRow, int mCol) {
    for (auto it : rowMap[mRow]) {
        if (holes[it].col == mCol) {
            holes[it].check = false;
            break;
        }
    }
}

int hitMarble(int mRow, int mCol, int mK) {
    step++;
    int resHole = -1; // Initialize to -1 to indicate no hole found
    while (mK--) {
        int dist = 9999999;

        for (auto it : rowMap[mRow]) {
            if (abs(holes[it].col - mCol) * 10 <= dist && holes[it].check == true && holes[it].cnt < step) {
                if (abs(holes[it].col - mCol) * 10 < dist || holes[it].col < holes[resHole].col) {
                    dist = abs(holes[it].col - mCol) * 10;
                    resHole = it;
                }
            }
        }

        for (auto it : colMap[mCol]) {
            if (abs(holes[it].row - mRow) * 10 <= dist && holes[it].check == true && holes[it].cnt < step) {
                if (abs(holes[it].row - mRow) * 10 < dist || holes[it].row < holes[resHole].row) {
                    dist = abs(holes[it].row - mRow) * 10;
                    resHole = it;
                }
            }
        }

        if (mRow >= leftDiagonal[mRow + mCol].sr && mRow <= leftDiagonal[mRow + mCol].er) {
            for (auto it : leftDiag[mRow + mCol]) {
                if (holes[it].row >= leftDiagonal[mRow + mCol].sr && holes[it].row <= leftDiagonal[mRow + mCol].er) {
                    if (abs(holes[it].row - mRow) * 14 <= dist && holes[it].check == true && holes[it].cnt < step) {
                        if (abs(holes[it].row - mRow) * 14 < dist || (holes[it].row == holes[resHole].row && holes[it].col < holes[resHole].col)) {
                            dist = abs(holes[it].row - mRow) * 14;
                            resHole = it;
                        }
                    }
                }
            }
        }

        if (mRow >= rightDiagonal[mRow - mCol + n].sr && mRow <= rightDiagonal[mRow - mCol + n].er) {
            for (auto it : rightDiag[mRow - mCol + n]) {
                if (holes[it].row >= rightDiagonal[mRow - mCol + n].sr && holes[it].row <= rightDiagonal[mRow - mCol + n].er) {
                    if (abs(holes[it].row - mRow) * 14 < dist && holes[it].check == true && holes[it].cnt < step) {
                        if (holes[it].row < holes[resHole].row || holes[it].col < holes[resHole].col) {
                            dist = abs(holes[it].row - mRow) * 14;
                            resHole = it;
                        }
                    }
                }
            }
        }

        if (dist == 9999999) break;
        else if (resHole != -1) {
            mRow = holes[resHole].row;
            mCol = holes[resHole].col;
            holes[resHole].cnt = step;
        }
    }

    if (resHole == -1) return -1;
    return resHole;
}