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 = mBRow;
		}
		else {
			leftDiagonal[mARow + mACol].sr = mBRow;
			leftDiagonal[mARow + mACol].sc = mBCol;
			leftDiagonal[mARow + mACol].er = mARow;
			leftDiagonal[mARow + mACol].ec = mARow;
		}
	} 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 = 0;
	while(mK--) {
		int dist = 9999999;
		for(auto it:rowMap[mRow]) {
			if(abs(holes[it].col - mCol) * 10 <= dist && holes[it].check == true 
				&& holes[it].col < holes[resHole].col && holes[it].cnt < step) 
			{
				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].row < holes[resHole].row && holes[it].cnt < step) 
			{
				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
						&& (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
						&& (holes[it].row < holes[resHole].row || holes[it].col < holes[resHole].col))
					{
						dist = abs(holes[it].row - mRow) * 14;
						resHole = it;
					}
				}
			}
		}
		mRow = holes[resHole].row;
		mCol = holes[resHole].col;
		holes[resHole].cnt = step;
	}
	return resHole;
}