Submission #235767

#	Time	Username	Problem	Language	Result	Execution time	Memory
235767		FlowerOfSorrow	Split the sequence (APIO14_sequence)	C++17	100 / 100	374 ms	6460 KiB

sequence

#include "bits/stdc++.h"
#include "ext/pb_ds/assoc_container.hpp"
#include "ext/pb_ds/tree_policy.hpp"
#include "ext/rope"
using namespace std;
using namespace chrono;
using namespace __gnu_pbds;
using namespace __gnu_cxx;
mt19937 rng(high_resolution_clock::now().time_since_epoch().count());
mt19937_64 rngll(high_resolution_clock::now().time_since_epoch().count());
#define lambdify(x) [&](auto &&...args){ return x(forward<decltype(args)>(args)...); }
template<typename T, typename U> T &ctmax(T &x, const U &y){ return x = max<T>(x, y); }
template<typename T, typename U> T &ctmin(T &x, const U &y){ return x = min<T>(x, y); }
template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

struct line{
	long long d, k, p;
	int ind;
	long long eval(long long x){ return d * x + k; }
};
template<bool GET_MAX = true>
struct sorted_line_container: deque<line>{
	static constexpr long long inf = numeric_limits<long long>::max();
	// (for doubles, use inf = 1/.0, div(a,b) = a/b)
	long long div(long long a, long long b){ return a / b - ((a ^ b) < 0 && a % b); }
	bool isect_front(iterator x, iterator y){
		if(y == end()){ x->p = inf; return false; }
		else{ x->p = div(y->k - x->k, x->d - y->d); return x->p >= y->p; }
	}
	bool isect_back(reverse_iterator x, reverse_iterator y){
		if(x == rend()) return false;
		else{ x->p = div(y->k - x->k, x->d - y->d); return x->p >= y->p; }
	}
	void push(line L){
		if(!GET_MAX) L.d = -L.d, L.k = -L.k;
		if(empty() || L.d < front().d){
			L.p = 0, push_front(L), isect_front(begin(), ++ begin());
			while(int(size()) >= 2 && isect_front(begin(), ++ begin())) erase(++ begin());
		}
		else if(L.d > back().d){
			L.p = inf, push_back(L); isect_back(++ rbegin(), rbegin());
			while(int(size()) >= 2 && isect_back(++ ++ rbegin(), ++ rbegin())) erase(-- -- end()), isect_back(++ rbegin(), rbegin());
		}
		else assert(false);
	}
	long long dec_query(long long x){
		while(int(size()) >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
		return rbegin()->eval(x) * (GET_MAX ? 1 : -1);
	}
	long long inc_query(long long x){
		while(int(size()) >= 2 && begin()->eval(x) <= (++ begin())->eval(x)) pop_front();
		return begin()->eval(x) * (GET_MAX ? 1 : -1);
	}
	long long query(long long x){
		if(int(size()) == 1) return begin()->eval(x) * (GET_MAX ? 1 : -1);
		int low = 0, high = int(size()) - 1;
		if(begin()->eval(x) >= (++ begin())->eval(x)) return begin()->eval(x) * (GET_MAX ? 1 : -1);
		while(high - low > 1){
			int mid = low + high >> 1;
			(*this)[mid].eval(x) < (*this)[mid + 1].eval(x) ? low = mid : high = mid;
		}
		return (*this)[low + 1].eval(x) * (GET_MAX ? 1 : -1);
	}
};

template<typename Pred>
long long custom_binary_search(long long low0, long long high0, const long long &step, Pred p, const bool &is_left = true){
	auto low = low0 / step - ((low0 ^ step) < 0 && low0 % step), high = high0 / step + ((high0 ^ step) > 0 && high0 % step);
	const auto rem = low0 - low * step;
	while(high - low > 1){
		auto mid = low + (high - low >> 1);
		(p(mid * step + rem) == is_left ? low : high) = mid;
	}
	return (is_left ? low : high) * step + rem;
}
 
template<typename DP, bool GET_MAX = true>
pair<long long, vector<int>> LagrangeDP(int n, DP f, long long k, long long low, long long high){
	long long resp, resq;
	vector<int> prevp(n + 1), cntp(n + 1), prevq(n + 1), cntq(n + 1);
	auto pred = [&](long long lambda){
		swap(resp, resq), swap(prevp, prevq), swap(cntp, cntq);
		resp = f(lambda, prevp, cntp);
		return GET_MAX ? cntp.back() <= k : cntp.back() >= k;
	};
	long long lambda = custom_binary_search(2 * low - 1, 2 * high + 1, 2, pred);
	pred(lambda + 2), pred(lambda);
	if(cntp.back() == k){
		vector<int> path{n};
		for(int u = n; u; ) path.push_back(u = prevp[u]);
		return {resp - lambda * k >> 1, path};
	}
	else{
		resp = resp - lambda * cntp.back() >> 1, resq = resq - (lambda + 2) * cntq.back() >> 1;
		long long res = resp + (resq - resp) / (cntq.back() - cntp.back()) * (k - cntp.back());
		int i = n, j = n, d = k - cntp.back();
		while(1){
			if(prevp[i] <= prevq[j]){
				while(prevp[i] <= prevq[j] && cntq[j] - cntp[i] > d) j = prevq[j];
				if(prevp[i] <= prevq[j] && cntq[j] - cntp[i] == d) break;
			}
			else i = prevp[i], j = prevq[j];
		}
		vector<int> path{n};
		for(int u = n; u != i; ) path.push_back(u = prevp[u]);
		path.push_back(prevq[j]);
		for(int u = prevq[j]; u; ) path.push_back(u = prevq[u]);
		return {res, path};
	}
}
 
int main(){
	cin.tie(0)->sync_with_stdio(0);
	int n, k;
	cin >> n >> k, ++ k;
	vector<pair<long long, int>> a(n);
	for(auto &p: a){
		static int cnt = 1;
		cin >> p.first, p.second = cnt ++;
	}
	int pcnt = 0;
	for(int i = 0; i < n; ++ i){
		if(a[i].first){
			++ pcnt;
		}
	}
	if(pcnt <= k){
		long long sum = 0, sqsum = 0;
		for(auto &[l, r]: a){
			sum += l, sqsum += l * l;
		}
		sum = sum * sum - sqsum >> 1;
		vector<int> flag(n - 1);
		int cur = 1;
		for(int i = 0; i < n - 1 && cur < k; ++ i){
			if(a[i].first){
				flag[i] = true;
				++ cur;
			}
		}
		for(int i = 0; i < n - 1 && cur < k; ++ i){
			if(!flag[i]){
				flag[i] = true;
				++ cur;
			}
		}
		cout << sum << "\n";
		for(int i = 0; i < n - 1; ++ i){
			if(flag[i]){
				cout << i + 1 << " ";
			}
		}
		return 0;
	}
	a.resize(stable_partition(a.begin(), a.end(), [](pair<long long, int> x){ return x.first; }) - a.begin());
	n = int(a.size());
	vector<long long> SUM(n + 1);
	for(int i = 0; i < n; ++ i){
		SUM[i + 1] = SUM[i] + a[i].first;
	}
	auto solve = [&](const long long &lambda, vector<int> &prev, vector<int> &cnt){
		long long res;
		sorted_line_container lc;
		lc.push({0, lambda, 0, 0});
		for(int i = 1; i <= n; ++ i){
			lc.push({2 * SUM[i], (res = lc.inc_query(SUM[i])) - 2 * SUM[i] * SUM[i] + lambda, 0, i});
			prev[i] = lc.front().ind;
			cnt[i] = cnt[prev[i]] + 1;
		}
		return res;
	};
	auto [res, seq] = LagrangeDP(n, solve, k, -1e18, 0);
	cout << res << "\n";
	for(int i = 1; i < int(seq.size()) - 1; ++ i){
		cout << a[seq[i] - 1].second << " ";
	}
	return 0;
}

/*

*/

////////////////////////////////////////////////////////////////////////////////////////
//                                                                                    //
//                                   Coded by Aeren                                   //
//                                                                                    //
////////////////////////////////////////////////////////////////////////////////////////

Compilation message (stderr)

sequence.cpp: In member function 'long long int sorted_line_container<GET_MAX>::dec_query(long long int)':
sequence.cpp:47:3: warning: this 'while' clause does not guard... [-Wmisleading-indentation]
   while(int(size()) >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
   ^~~~~
sequence.cpp:47:86: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'while'
   while(int(size()) >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
                                                                                      ^~~~~~
sequence.cpp: In function 'int main()':
sequence.cpp:129:18: warning: unused variable 'r' [-Wunused-variable]
   for(auto &[l, r]: a){
                  ^
sequence.cpp:132:19: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   sum = sum * sum - sqsum >> 1;
         ~~~~~~~~~~^~~~~~~
sequence.cpp: In instantiation of 'std::pair<long long int, std::vector<int> > LagrangeDP(int, DP, long long int, long long int, long long int) [with DP = main()::<lambda(const long long int&, std::vector<int>&, std::vector<int>&)>; bool GET_MAX = true]':
sequence.cpp:172:52:   required from here
sequence.cpp:91:16: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   return {resp - lambda * k >> 1, path};
           ~~~~~^~~~~~~~~~~~
sequence.cpp:94:15: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   resp = resp - lambda * cntp.back() >> 1, resq = resq - (lambda + 2) * cntq.back() >> 1;
          ~~~~~^~~~~~~~~~~~~~~~~~~~~~
sequence.cpp:94:56: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   resp = resp - lambda * cntp.back() >> 1, resq = resq - (lambda + 2) * cntq.back() >> 1;
                                                   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~
sequence.cpp: In instantiation of 'long long int custom_binary_search(long long int, long long int, const long long int&, Pred, const bool&) [with Pred = LagrangeDP(int, DP, long long int, long long int, long long int) [with DP = main()::<lambda(const long long int&, std::vector<int>&, std::vector<int>&)>; bool GET_MAX = true]::<lambda(long long int)>]':
sequence.cpp:86:41:   required from 'std::pair<long long int, std::vector<int> > LagrangeDP(int, DP, long long int, long long int, long long int) [with DP = main()::<lambda(const long long int&, std::vector<int>&, std::vector<int>&)>; bool GET_MAX = true]'
sequence.cpp:172:52:   required from here
sequence.cpp:71:26: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   auto mid = low + (high - low >> 1);
                     ~~~~~^~~~~
sequence.cpp: In lambda function:
sequence.cpp:83:3: warning: 'res' may be used uninitialized in this function [-Wmaybe-uninitialized]
   resp = f(lambda, prevp, cntp);
   ^~~~
sequence.cpp: In function 'int main()':
sequence.cpp:83:3: warning: 'res' may be used uninitialized in this function [-Wmaybe-uninitialized]
   resp = f(lambda, prevp, cntp);
   ^~~~
sequence.cpp:162:13: note: 'res' was declared here
   long long res;
             ^~~

• Subtask 1: 11 / 11
• Subtask 2: 11 / 11
• Subtask 3: 11 / 11
• Subtask 4: 17 / 17
• Subtask 5: 21 / 21
• Subtask 6: 29 / 29