Submission #199301

#	Time	Username	Problem	Language	Result	Execution time	Memory
199301		Mahotsukai	Split the sequence (APIO14_sequence)	C++17	100 / 100	379 ms	5948 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 __gnu_pbds;
using namespace __gnu_cxx;
template<class L, class R> istream &operator>>(istream &in, pair<L, R> &p){ return in >> p.first >> p.second; }
template<class Tuple, size_t ...Is> void read_tuple(istream &in, Tuple &t, index_sequence<Is...>){ ((in >> get<Is>(t)), ...); }
template<class ...Args> istream &operator>>(istream &in, tuple<Args...> &t){ read_tuple(in, t, index_sequence_for<Args...>{}); return in; }
template<class ...Args, template<class...> class T> istream &operator>>(enable_if_t<!is_same_v<T<Args...>, string>, istream> &in, T<Args...> &arr){ for(auto &x: arr) in >> x; return in; }
template<class L, class R> ostream &operator<<(ostream &out, const pair<L, R> &p){ return out << "(" << p.first << ", " << p.second << ")"; }
//template<class L, class R> ostream &operator<<(ostream &out, const pair<L, R> &p){ return out << p.first << " " << p.second << "\n"; }
template<class Tuple, size_t ...Is> void print_tuple(ostream &out, const Tuple &t, index_sequence<Is...>){ ((out << (Is ? " " : "") << get<Is>(t)), ...); }
template<class ...Args> ostream &operator<<(ostream &out, const tuple<Args...> &t){ print_tuple(out, t, index_sequence_for<Args...>{}); return out << "\n"; }
template<class ...Args, template<class...> class T> ostream &operator<<(enable_if_t<!is_same_v<T<Args...>, string>, ostream> &out, const T<Args...> &arr){ for(auto &x: arr) out << x << " "; return out << "\n"; }
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
mt19937_64 rngll(chrono::steady_clock::now().time_since_epoch().count());
#define all(a) a.begin(), a.end()
#define sz(a) (int)a.size()
typedef long long ll;
typedef vector<int> vi; typedef vector<ll> vl; typedef vector<double> vd; typedef vector<string> vs;
typedef pair<int, int> pii; typedef pair<ll, ll> pll; typedef pair<int, ll> pil; typedef pair<ll, int> pli;
typedef vector<pii> vpii; typedef vector<pil> vpil; typedef vector<pli> vpli; typedef vector<pll> vpll;
template<class T> T ctmax(T &x, const T &y){ return x = max(x, y); }
template<class T> T ctmin(T &x, const T &y){ return x = min(x, y); }
template<class T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef tuple<int, int, int> tpl; typedef vector<tpl> vtpl;

struct line{
	ll d, k, p;
	int ind;
	ll eval(ll x){ return d * x + k; }
};
template<bool GET_MAX = true>
struct sorted_line_container: deque<line>{
	// (for doubles, use inf = 1/.0, div(a,b) = a/b)
	const ll inf = LLONG_MAX;
	ll div(ll a, ll b){ return a / b - ((a ^ b) < 0 && a % b); }
	bool isect_front(iterator x, iterator y){
		if(y == this->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 == this->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(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(size() >= 2 && isect_back(++ ++ rbegin(), ++ rbegin())) erase(-- -- end()), isect_back(++ rbegin(), rbegin());
		}
		else assert(false);
	}
	ll dec_query(ll x){
		while(size() >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
		return rbegin()->eval(x) * (GET_MAX ? 1 : -1);
	}
	ll inc_query(ll x){
		while(size() >= 2 && begin()->eval(x) <= (++ begin())->eval(x)) pop_front();
		return begin()->eval(x) * (GET_MAX ? 1 : -1);
	}
	ll query(ll x){
		if(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<class Pred>
ll custom_binary_search(ll low, ll high, const ll &step, Pred p, bool is_left = true){
	assert(low < high && (high - low) % step == 0);
	const ll rem = low % step;
	if(is_left){
		while(high - low > step){
			ll mid = low + (high - low >> 1);
			mid = mid / step * step + rem;
			p(mid) ? low = mid : high = mid;
		}
		return low;
	}
	else{
		while(high - low > step){
			ll mid = low + (high - low >> 1);
			mid = mid / step * step + rem;
			p(mid) ? high = mid : low = mid;
		}
		return high;
	}
}

template<class DP, bool GET_MAX = true>
pair<ll, vi> LagrangeDP(int n, DP f, ll k, ll low, ll high){
	ll resp, resq;
	vi prevp(n + 1), cntp(n + 1), prevq(n + 1), cntq(n + 1);
	auto pred = [&](ll lambda){
		swap(resp, resq), swap(prevp, prevq), swap(cntp, cntq);
		resp = f(lambda, prevp, cntp);
		return GET_MAX ? cntp.back() <= k : cntp.back() >= k;
	};
	ll lambda = custom_binary_search(2 * low - 1, 2 * high + 1, 2, pred);
	pred(lambda + 2), pred(lambda);
	if(cntp.back() == k){
		vi 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;
		ll 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];
		}
		vi 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;
	vpli 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){
		ll sum = 0, sqsum = 0;
		for(auto &[l, r]: a){
			sum += l, sqsum += l * l;
		}
		sum = sum * sum - sqsum >> 1;
		vi 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(all(a), [](pli x){ return x.first; }) - a.begin());
	n = sz(a);
	vl SUM(n + 1);
	for(int i = 0; i < n; ++ i){
		SUM[i + 1] = SUM[i] + a[i].first;
	}
	auto solve = [&](const ll &lambda, vi &prev, vi &cnt){
		ll 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 < sz(seq) - 1; ++ i){
		cout << a[seq[i] - 1].second << " ";
	}
	return 0;
}

/* 

7 3
4 1 3 4 0 2 3

14 7
3 2 1 4 1 3 3 3 2 1 1 1 4 2 

p: 14->12->9 ->7->6->4->(2->0)
q: 14->13->12->9->7->6->4->3->(1->0)

dp[i][j]: partition i times, max cost for first j

dp[i][j] = max{k < j}( dp[i-1][k] + cost[k][j] )
cost[k][j] = sum(0, k) * sum(k, j) = K * (J - K) = K * J - K^2
a<=b<=c<=d
cost[a][c] + cost[b][d] = A(C-A) + B(D-B)
cost[a][d] + cost[b][c] = A(D-A) + B(C-B)

dp'[i] = max{j < i}( dp[j] - SUM^2[j] + SUM[j] * SUM[i] + lambda )
*/

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

Compilation message (stderr)

sequence.cpp: In member function 'll sorted_line_container<GET_MAX>::dec_query(ll)':
sequence.cpp:61:3: warning: this 'while' clause does not guard... [-Wmisleading-indentation]
   while(size() >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
   ^~~~~
sequence.cpp:61:81: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'while'
   while(size() >= 2 && rbegin()->eval(x) <= (++ rbegin())->eval(x)) pop_back(); rbegin()->p = inf;
                                                                                 ^~~~~~
sequence.cpp: In function 'int main()':
sequence.cpp:154:18: warning: unused variable 'r' [-Wunused-variable]
   for(auto &[l, r]: a){
                  ^
sequence.cpp:157: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, ll, ll, ll) [with DP = main()::<lambda(const ll&, vi&, vi&)>; bool GET_MAX = true; ll = long long int]':
sequence.cpp:197:52:   required from here
sequence.cpp:116:16: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   return {resp - lambda * k >> 1, path};
           ~~~~~^~~~~~~~~~~~
sequence.cpp:119:15: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   resp = resp - lambda * cntp.back() >> 1, resq = resq - (lambda + 2) * cntq.back() >> 1;
          ~~~~~^~~~~~~~~~~~~~~~~~~~~~
sequence.cpp:119: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 'll custom_binary_search(ll, ll, const ll&, Pred, bool) [with Pred = LagrangeDP(int, DP, ll, ll, ll) [with DP = main()::<lambda(const ll&, vi&, vi&)>; bool GET_MAX = true; ll = long long int]::<lambda(ll)>; ll = long long int]':
sequence.cpp:111:34:   required from 'std::pair<long long int, std::vector<int> > LagrangeDP(int, DP, ll, ll, ll) [with DP = main()::<lambda(const ll&, vi&, vi&)>; bool GET_MAX = true; ll = long long int]'
sequence.cpp:197:52:   required from here
sequence.cpp:86:25: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
    ll mid = low + (high - low >> 1);
                    ~~~~~^~~~~
sequence.cpp:94:25: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
    ll mid = low + (high - low >> 1);
                    ~~~~~^~~~~
sequence.cpp: In lambda function:
sequence.cpp:108:3: warning: 'res' may be used uninitialized in this function [-Wmaybe-uninitialized]
   resp = f(lambda, prevp, cntp);
   ^~~~

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