test/slope_trick.test.cpp

test/slope_trick.test.cpp · GitHub で見る

依存ライブラリ

コード

#define PROBLEM "https://yukicoder.me/problems/no/2304"
#include "../data_structure/slope_trick.hpp"
#include <algorithm>
#include <cstdint>
#include <iostream>
#include <vector>

int main() {
	std::cin.tie(nullptr)->sync_with_stdio(false);
	int n;
	std::cin >> n;
	std::vector<std::int64_t> a(n);
	for (auto &x : a) std::cin >> x;
	std::sort(a.begin(), a.end());
	SlopeTrick<> slope;
	slope.reserve(n);
	for (auto x : a) {
		slope.shift(1);
		slope.clear_right();
		slope.add_abs(x);
	}
	std::cout << slope.min() << '\n';
}

#line 1 "test/slope_trick.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/2304"
#line 1 "data_structure/slope_trick.hpp"



#line 1 "components/d_ary_heap.hpp"



#include <algorithm>
#include <cassert>
#include <cstddef>
#include <functional>
#include <utility>
#include <vector>

template <class T, class Compare = std::less<T>, std::size_t D = 2>
class DaryHeap {
	static_assert(D >= 2);
	std::vector<T> data;
	[[no_unique_address]] Compare compare{};
	void up(std::size_t i, T value) {
		while (i) {
			std::size_t p = (i - 1) / D;
			if (!compare(data[p], value)) break;
			data[i] = std::move(data[p]);
			i = p;
		}
		data[i] = std::move(value);
	}
	void down(std::size_t i, T value) {
		std::size_t n = data.size();
		while (i < n / D + (n % D > 1)) {
			std::size_t best = i * D + 1, end = n;
			if constexpr (D <= std::size_t(-1) / 2) end = std::min(best + D, n);
			for (std::size_t j = best + 1; j < end; ++j)
				if (compare(data[best], data[j])) best = j;
			if (!compare(value, data[best])) break;
			data[i] = std::move(data[best]);
			i = best;
		}
		data[i] = std::move(value);
	}

  public:
	DaryHeap() = default;
	explicit DaryHeap(Compare compare) : compare(std::move(compare)) {}
	explicit DaryHeap(std::vector<T> values, Compare compare = {}) : data(std::move(values)), compare(std::move(compare)) {
		for (std::size_t i = data.size() / D + (data.size() % D > 1); i--;) down(i, std::move(data[i]));
	}
	bool empty() const { return data.empty(); }
	std::size_t size() const { return data.size(); }
	void reserve(std::size_t n) { data.reserve(n); }
	void clear() { data.clear(); }
	const std::vector<T> &values() const { return data; }
	decltype(auto) top() const {
		assert(!empty());
		return data.front();
	}
	void push(T value) {
		data.push_back(std::move(value));
		up(data.size() - 1, std::move(data.back()));
	}
	T replace_top(T value) {
		assert(!empty());
		T result = std::move(data.front());
		down(0, std::move(value));
		return result;
	}
	T push_pop(T value) {
		if (empty() || !compare(value, data.front())) return value;
		return replace_top(std::move(value));
	}
	T pop_value() {
		assert(!empty());
		T result = std::move(data.front());
		if (data.size() == 1) data.pop_back();
		else {
			T value = std::move(data.back());
			data.pop_back();
			down(0, std::move(value));
		}
		return result;
	}
	void pop() { pop_value(); }
	void swap(DaryHeap &other) noexcept(noexcept(std::swap(compare, other.compare))) {
		data.swap(other.data);
		std::swap(compare, other.compare);
	}
};


#line 7 "data_structure/slope_trick.hpp"
#include <cstdint>
#line 9 "data_structure/slope_trick.hpp"
#include <optional>
#include <type_traits>
#line 12 "data_structure/slope_trick.hpp"

template <class T = std::int64_t, class Sum = T>
class SlopeTrick {
	static_assert((std::is_arithmetic_v<T> && std::is_signed_v<T>) || std::is_same_v<T, __int128_t>);
	static_assert((std::is_arithmetic_v<Sum> && std::is_signed_v<Sum>) || std::is_same_v<Sum, __int128_t>);
	DaryHeap<T> left;
	DaryHeap<T, std::greater<T>> right;
	T offset_left = 0, offset_right = 0;
	Sum minimum = 0;

  public:
	SlopeTrick() = default;
	SlopeTrick(const SlopeTrick &) = default;
	SlopeTrick &operator=(const SlopeTrick &) = default;
	SlopeTrick(SlopeTrick &&other) noexcept : left(std::move(other.left)), right(std::move(other.right)), offset_left(std::exchange(other.offset_left, 0)), offset_right(std::exchange(other.offset_right, 0)), minimum(std::exchange(other.minimum, Sum(0))) {}
	SlopeTrick &operator=(SlopeTrick &&other) noexcept {
		if (this != &other) {
			left = std::move(other.left);
			right = std::move(other.right);
			offset_left = std::exchange(other.offset_left, 0);
			offset_right = std::exchange(other.offset_right, 0);
			minimum = std::exchange(other.minimum, Sum(0));
		}
		return *this;
	}
	std::size_t size() const { return left.size() + right.size(); }
	bool empty() const { return size() == 0; }
	void reserve(std::size_t n) {
		left.reserve(n);
		right.reserve(n);
	}
	void clear_left() {
		left.clear();
		offset_left = 0;
	}
	void clear_right() {
		right.clear();
		offset_right = 0;
	}
	void clear() {
		clear_left();
		clear_right();
		minimum = 0;
	}
	Sum min() const { return minimum; }
	std::pair<std::optional<T>, std::optional<T>> argmin() const {
		return {left.empty() ? std::nullopt : std::optional<T>(left.top() + offset_left), right.empty() ? std::nullopt : std::optional<T>(right.top() + offset_right)};
	}
	void add_all(Sum a) { minimum += a; }
	void add_a_minus_x(T a) {
		if (!right.empty() && right.top() + offset_right < a) {
			T r = right.replace_top(a - offset_right) + offset_right;
			minimum += Sum(a) - Sum(r);
			left.push(r - offset_left);
		} else left.push(a - offset_left);
	}
	void add_x_minus_a(T a) {
		if (!left.empty() && a < left.top() + offset_left) {
			T l = left.replace_top(a - offset_left) + offset_left;
			minimum += Sum(l) - Sum(a);
			right.push(l - offset_right);
		} else right.push(a - offset_right);
	}
	void add_abs(T a) {
		if (!left.empty() && a < left.top() + offset_left) {
			T l = left.replace_top(a - offset_left) + offset_left;
			minimum += Sum(l) - Sum(a);
			left.push(a - offset_left);
			right.push(l - offset_right);
		} else if (!right.empty() && right.top() + offset_right < a) {
			T r = right.replace_top(a - offset_right) + offset_right;
			minimum += Sum(a) - Sum(r);
			right.push(a - offset_right);
			left.push(r - offset_left);
		} else left.push(a - offset_left), right.push(a - offset_right);
	}
	void shift(T a, T b) {
		assert(a <= b);
		offset_left += a;
		offset_right += b;
	}
	void shift(T a) { shift(a, a); }
	Sum eval(T x) const {
		Sum value = minimum;
		for (T a : left.values())
			if (x < a + offset_left) value += Sum(a + offset_left) - Sum(x);
		for (T a : right.values())
			if (a + offset_right < x) value += Sum(x) - Sum(a + offset_right);
		return value;
	}
	void swap(SlopeTrick &other) noexcept {
		left.swap(other.left);
		right.swap(other.right);
		std::swap(offset_left, other.offset_left);
		std::swap(offset_right, other.offset_right);
		std::swap(minimum, other.minimum);
	}
	void merge(SlopeTrick &other) {
		if (this == &other) {
			SlopeTrick copy(other);
			merge(copy);
			return;
		}
		if (size() < other.size()) swap(other);
		minimum += other.minimum;
		for (T a : other.left.values()) add_a_minus_x(a + other.offset_left);
		for (T a : other.right.values()) add_x_minus_a(a + other.offset_right);
		other = SlopeTrick{};
	}
	void merge(SlopeTrick &&other) { merge(other); }
};


#line 5 "test/slope_trick.test.cpp"
#include <iostream>
#line 7 "test/slope_trick.test.cpp"

int main() {
	std::cin.tie(nullptr)->sync_with_stdio(false);
	int n;
	std::cin >> n;
	std::vector<std::int64_t> a(n);
	for (auto &x : a) std::cin >> x;
	std::sort(a.begin(), a.end());
	SlopeTrick<> slope;
	slope.reserve(n);
	for (auto x : a) {
		slope.shift(1);
		slope.clear_right();
		slope.add_abs(x);
	}
	std::cout << slope.min() << '\n';
}