Untitled

#include "judge.h"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <set>
#include <vector>
using namespace std;

constexpr int N = 5e4 + 7;
constexpr int T = 50 + 7;
set<pair<int, int>> cord_y_points[N][T];

#ifndef PEKI_DEBUG
	#define PEKI_DEBUG 0
#endif

namespace {
	int castToIntWithEpsilon(double v) {
		constexpr double eps = 1e-5;
		if(static_cast<int>(v + eps) > static_cast<int>(v)){
			return static_cast<int>(v) + 1;
		}
		return static_cast<int>(v);
	}
}

struct Point {
	int x;
	int y;
	int id;
	int type;
	Point(int _x, int _y, int _id, int _type) : x(_x), y(_y), id(_id), type(_type) {

	}
	Point() = default;
};

class circle : public icircle {
public:
	circle() {
		points.resize(N);
	}
	virtual void addGemstone(int x, int y, int id, int type);
	virtual void removeGemstone(int id);
	virtual int countGemstones(int x, int y, int r, int type);
	virtual int findClosestGemstones(int x, int y, int r, int type, int res[3]);
private:
	long long dist(const Point &p1, const Point &p2);
	bool checkDist(const Point &p1, const Point &p2, int r);
	int countForType(int x, int y, int r, int type);

	pair <int, int> getInterval(long long a, long long b, long long c);

	vector<Point> points;
};

long long circle::dist(const Point &p1, const Point &p2) {
	return static_cast<long long>(p1.x - p2.x) * (p1.x - p2.x) + static_cast<long long>(p1.y - p2.y) * (p1.y - p2.y);
}

bool circle::checkDist(const Point &p1, const Point &p2, int r) {
	auto distance = dist(p1, p2);
	return distance <= static_cast<long long>(r) * r;
}

pair <int, int> circle::getInterval(long long a, long long b, long long c)
{
	double delta = b * b - 4 * a * c;
	// for given internal should always exists solution
	assert(delta >= 0);
	double sqrt_delta = sqrt(delta);
	
	auto x1 = static_cast<double>(-b - sqrt_delta) / 2 * a;
	auto x2 = static_cast<double>(-b + sqrt_delta) / 2 * a;
	
	if(x1 < 0) { x1 = 0.0; }
	if(x2 < 0) { x2 = 0.0; }
	if(x1 > x2) swap(x1, x2);

	return {castToIntWithEpsilon(x1), castToIntWithEpsilon(x2)};
}

int circle::countForType(int x, int y, int r, int type)
{
	int y_a = max(0, y - r), y_b = max(0, y + r);
	int result = 0;
	for(int y_i = y_a; y_i <= y_b; y_i++) {
		const auto& x_points = cord_y_points[y_i][type];
		if (x_points.empty()) { continue; }
		auto c = static_cast<long long>(y_i - y) * (y_i - y)
						 + static_cast<long long>(x) * x
						 - static_cast<long long>(r) * r;
		auto interval = getInterval(1.0, -2 * x, c);
#if PEKI_DEBUG
		if(x_points.size() > 0) {
			std::cout << "(x: " << x << ", y: " << y << " r: " << r << ", yi: " << y_i << ") -> {"
			<< interval.first << ", " << interval.second << "}" << std::endl; 
		}
		for(auto &p : x_points) {
			std::cout << "(xi: " << p.first << ", id: " << p.second << ")" << std::endl;
		}
#endif
		auto a_it = x_points.lower_bound({interval.first, 0});
		if (a_it == x_points.end()) { continue; }
		auto b_it = x_points.upper_bound({interval.second, N});
		result += distance(a_it, b_it);
	}
	return result;
}

void circle::addGemstone(int x, int y, int id, int type)
{
	points[id] = {x, y, id, type};
	cord_y_points[y][type].emplace(x, id);
}

void circle::removeGemstone(int id)
{
	cord_y_points[points[id].y][points[id].type].erase({points[id].x, id});
}

int circle::countGemstones(int x, int y, int r, int type) 
{
	int result = 0;
	if (type == 0) {
		for(int i = 0; i < T; i++){
			result += countForType(x, y, r, i);
		}
	}
	else {
		result += countForType(x, y, r, type);
	}
	return result;
}

int circle::findClosestGemstones(int x, int y, int r, int type, int res[3])
{
	set <pair <long long, int> > result_set;
	vector <pair <long long, int> > result_points;
	int y_a = max(0, y - r), y_b = max(0, y + r);
	for(int y_i = y_a; y_i <= y_b; y_i++) {
		const auto& x_points = cord_y_points[y_i][type];
		if (x_points.empty()) { continue; }
		auto bound_it = x_points.lower_bound({x, 0});
		if (bound_it == x_points.end()) {
			int ind = 0;
			for(auto it = x_points.rbegin(); it != x_points.rend() && ind < 3; it++) {
				result_points.emplace_back(dist({x, y, 0, 0}, {it->first, y_i, 0, 0}), it->second);
				ind++;
			}
		}
		else {
			auto left = bound_it;
			auto right = bound_it; right++;
			bool found = false;
			while(!found) {
				if(result_points.size() == 3) {
					found = true;
					continue;
				}
				bool left_ok = true;
				bool right_ok = true;
				if(left == x_points.end()) {
					left_ok = false;
				}
				if(right == x_points.end()){
					right_ok = false;
				}
				if(!left_ok && !right_ok) {
					found = true;
					continue;
				}
				if(right_ok && left_ok) {
					auto left_distance = dist({points[left->second].x, points[left->second].y, 0, 0}, {x, y, 0, 0});
					auto right_distance = dist({points[right->second].x, points[right->second].y, 0, 0}, {x, y, 0, 0});
					if(left_distance < right_distance) {
						result_points.emplace_back(left_distance, left->second);
						if(left == x_points.begin()) { left = x_points.end(); }
						else { left--; }
					} else {
						result_points.emplace_back(right_distance, right->second);
						right++;
					}
				} else if(right_ok && !left_ok)
				{
					auto right_distance = dist({points[right->second].x, points[right->second].y, 0, 0}, {x, y, 0, 0});
					result_points.emplace_back(right_distance, right->second);
					right++;
				} else {
					auto left_distance = dist({points[left->second].x, points[left->second].y, 0, 0}, {x, y, 0, 0});
					result_points.emplace_back(left_distance, left->second);
					if(left == x_points.begin()) { left = x_points.end(); }
					else { left--; }
				}
			}
		}
		for(const auto &p : result_points) {
			if(p.first <= static_cast<long long>(r) * r) {
				result_set.emplace(p.first, p.second);
			}
		}
	}
	int ind = 0;
	for(auto it = result_set.begin(); it != result_set.end() && ind < 3; it++) {
		res[ind++] = it->second;
	}
	return ind;
}