# Untitled

unknown
plain_text
18 days ago
6.1 kB
2
Indexable
Never
```#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;
}```