Untitled

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

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] = { -1, -1, 0, false };
    }
    for (int i = 0; i < 6001; i++) {
        leftDiagonal[i] = { -1, -1, -1, -1 };
        rightDiagonal[i] = { -1, -1, -1, -1 };
    }
}

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

void addHole(int mRow, int mCol, int mID) {
    holes[mID] = { mRow, mCol, 0, 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].begin(); it != rowMap[mRow].end(); ++it) {
        if (holes[*it].col == mCol) {
            holes[*it].check = false;
            rowMap[mRow].erase(it);  // Erase from rowMap
            break;
        }
    }
}

int hitMarble(int mRow, int mCol, int mK) {
    step++;
    int resHole = -1;

    while (mK--) {
        int dist = 9999999;
        resHole = -1;  // Reset resHole for each iteration

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

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

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

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

        if (resHole == -1) break;  // No valid hole found, exit loop

        mRow = holes[resHole].row;
        mCol = holes[resHole].col;
        holes[resHole].cnt = step;  // Update count for this hole
    }

    return resHole;
}