#include "aliens.h"
#include <vector>
#include <algorithm>
#include <deque>
using namespace std;
typedef long long ll;
// represents a linear function for convex hull trick
class linearfunction {
private:
ll slope, intercept;
public:
linearfunction(ll m, ll c) : slope(m), intercept(c) {}
ll evaluate(ll x) const {
return slope * x + intercept;
}
ll getslope() const { return slope; }
ll getintercept() const { return intercept; }
};
// convex hull trick for maintaining minimum of linear functions
class convexhulltrick {
private:
deque<linearfunction> hull;
// check if middle line is redundant given left and right
bool isredundant(const linearfunction& left, const linearfunction& mid, const linearfunction& right) {
// using 128-bit integers to avoid overflow
__int128 lhs = (__int128)(left.getslope() - right.getslope()) *
(__int128)(mid.getintercept() - left.getintercept());
__int128 rhs = (__int128)(left.getslope() - mid.getslope()) *
(__int128)(right.getintercept() - left.getintercept());
return lhs <= rhs;
}
public:
void insert(linearfunction func) {
// handle duplicate slopes
if (!hull.empty() && hull.front().getslope() == func.getslope()) {
if (hull.front().getintercept() <= func.getintercept()) {
return;
}
hull.pop_front();
}
// maintain convex hull property
while (hull.size() >= 2) {
if (isredundant(func, hull[0], hull[1])) {
hull.pop_front();
} else {
break;
}
}
hull.push_front(func);
}
ll querymin(ll x) {
// remove suboptimal functions from the back
while (hull.size() >= 2) {
int lastidx = hull.size() - 1;
if (hull[lastidx].evaluate(x) >= hull[lastidx - 1].evaluate(x)) {
hull.pop_back();
} else {
break;
}
}
return hull.back().evaluate(x);
}
};
// segment representing a photo area
struct segment {
int start, end;
segment(int s, int e) : start(s), end(e) {}
bool operator<(const segment& other) const {
if (start != other.start) return start < other.start;
return end > other.end; // prefer longer segments for same start
}
ll coverage() const {
return (ll)(end - start + 1) * (end - start + 1);
}
};
// preprocess points to get minimal set of segments
vector<segment> preprocesssegments(int n, vector<int>& r, vector<int>& c) {
vector<segment> segments;
segments.reserve(n);
for (int i = 0; i < n; i++) {
int startpos = min(r[i], c[i]);
int endpos = max(r[i], c[i]);
segments.emplace_back(startpos, endpos);
}
sort(segments.begin(), segments.end());
// remove dominated segments
vector<segment> filtered;
int maxend = -1;
for (const auto& seg : segments) {
if (seg.end > maxend) {
filtered.push_back(seg);
maxend = seg.end;
}
}
return filtered;
}
// calculate overlap between consecutive segments
ll calculateoverlap(const segment& prev, const segment& curr) {
if (curr.start > prev.end) return 0;
ll overlap = prev.end - curr.start + 1;
return overlap * overlap;
}
// dp with lagrangian relaxation - returns (cost, count)
pair<ll, int> solvewithpenalty(const vector<segment>& segments, ll penalty) {
int n = segments.size();
vector<ll> dp(n + 1, 1e18);
vector<int> cnt(n + 1, 0);
dp[0] = 0;
cnt[0] = 0;
convexhulltrick cht;
for (int i = 0; i <= n; i++) {
if (i > 0) {
ll x = segments[i - 1].end;
ll mincost = cht.querymin(x);
dp[i] = mincost + x * x + penalty;
// find which transition gave us this value
for (int j = 0; j < i; j++) {
ll slope = -2 * segments[j].start + 2;
ll overlap = (j > 0) ? calculateoverlap(segments[j - 1], segments[j]) : 0;
ll base = segments[j].start - 1;
ll intercept = dp[j] + base * base - overlap;
ll transitioncost = slope * x + intercept + x * x + penalty;
if (abs(transitioncost - dp[i]) < 1e-9) {
cnt[i] = cnt[j] + 1;
break;
}
}
}
if (i < n) {
ll slope = -2 * segments[i].start + 2;
ll overlap = (i > 0) ? calculateoverlap(segments[i - 1], segments[i]) : 0;
ll base = segments[i].start - 1;
ll intercept = dp[i] + base * base - overlap;
cht.insert(linearfunction(slope, intercept));
}
}
return {dp[n] - penalty * cnt[n], cnt[n]};
}
long long take_photos(int n, int m, int k, vector<int> r, vector<int> c) {
auto segments = preprocesssegments(n, r, c);
n = segments.size();
if (n <= k) {
// can use one segment per point
ll total = 0;
for (int i = 0; i < n; i++) {
total += segments[i].coverage();
if (i > 0) {
total -= calculateoverlap(segments[i - 1], segments[i]);
}
}
return total;
}
// binary search on penalty using ternary search optimization
ll left = 0, right = (ll)m * m;
while (right - left > 2) {
ll mid1 = left + (right - left) / 3;
ll mid2 = right - (right - left) / 3;
auto [cost1, cnt1] = solvewithpenalty(segments, mid1);
auto [cost2, cnt2] = solvewithpenalty(segments, mid2);
ll adjusted1 = cost1 + mid1 * cnt1;
ll adjusted2 = cost2 + mid2 * cnt2;
if (adjusted1 > adjusted2) {
left = mid1;
} else {
right = mid2;
}
}
// check remaining candidates
ll result = 1e18;
for (ll penalty = left; penalty <= right; penalty++) {
auto [cost, numsegments] = solvewithpenalty(segments, penalty);
if (numsegments <= k) {
result = min(result, cost);
}
}
return result;
}
