Submission #199270

#	Time	Username	Problem	Language	Result	Execution time	Memory
199270		Mahotsukai	Split the sequence (APIO14_sequence)	C++17	100 / 100	376 ms	6768 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;

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;
	}
}

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);
	}
};

int main(){
	cin.tie(0)->sync_with_stdio(0);
	int n, k;
	cin >> n >> k, ++ k;
	vpli a(n);
	vl b;
	int N = n;
	for(auto &p: a){
		static int cnt = 1;
		cin >> p.first, p.second = cnt ++;
		b.push_back(p.first);
	}
	auto check_sol = [&](ll x, const vi &t){
		vl SUM(N + 1);
		partial_sum(all(b), SUM.begin() + 1);
		auto sum = [&](int i, int j){ return SUM[j] - SUM[i]; };
		ll res = 0;
		for(auto i: t){
			static int prev = 0;
			res += sum(prev, i) * sum(i, N);
			prev = i;
		}
		cout << "X = " << x << ", my seq implies = " << res << endl;
		assert(x == res);
	};
	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;
	}
	sorted_line_container lc;
	vi cnt(n + 1), prev(n + 1);
	bool FLAG = false;
	vi p, q, cntp, cntq;
	ll resp = 0, resq = 0;
	auto test = [&](ll lambda){
		swap(resp, resq);
		lc.clear();
		fill(all(cnt), 0);
		fill(all(prev), 0);
		lc.push({0, lambda, 0, 0});
		for(int i = 1; i <= n; ++ i){
			lc.push({2 * SUM[i], (resp = lc.inc_query(SUM[i])) - 2 * SUM[i] * SUM[i] + lambda, 0, i});
			prev[i] = lc.front().ind;
			cnt[i] = cnt[prev[i]] + 1;
		}
		if(FLAG){
			swap(p, q), p = prev;
			swap(cntp, cntq), cntp = cnt;
		}
		return cnt.back() <= k;
	};
	ll lambda = custom_binary_search(-(ll)2e18 - 1, 1, 2, test);
	FLAG = true;
	test(lambda + 2), test(lambda);
	if(cntp.back() == k){
		vi res;
		for(int u = n; p[u]; ){
			res.push_back(a[(u = p[u]) - 1].second);
		}
		reverse(all(res));
		cout << (resp - lambda * k) / 2 << "\n" << res;
	}
	else{
		resp = (resp - lambda * cntp.back()) / 2, resq = (resq - (lambda + 2) * cntq.back()) / 2;
		int i = n, j = n, d = k - cntp.back();
		while(1){ // prev[i] <= prev[j] < j <= i
			if(p[i] <= q[j]){
				while(p[i] <= q[j] && cntq[j] - cntp[i] > d) j = q[j];
				if(p[i] <= q[j] && cntq[j] - cntp[i] == d) break;
			}
			else i = p[i], j = q[j];
		}
		vi res;
		for(int u = n; u != i; ){
			res.push_back(a[(u = p[u]) - 1].second);
		}
		if(q[j]) res.push_back(a[q[j] - 1].second);
		for(int u = q[j]; q[u]; ){
			res.push_back(a[(u = q[u]) - 1].second);
		}
		reverse(all(res));
		vl SUM(N + 1);
		partial_sum(all(b), SUM.begin() + 1);
		auto sum = [&](int i, int j){ return SUM[j] - SUM[i]; };
		cout << resp + (resq - resp) / (cntq.back() - cntp.back()) * (k - cntp.back()) << "\n" << res;
	}
	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:83: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:83: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:135:18: warning: unused variable 'r' [-Wunused-variable]
   for(auto &[l, r]: a){
                  ^
sequence.cpp:138:19: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   sum = sum * sum - sqsum >> 1;
         ~~~~~~~~~~^~~~~~~
sequence.cpp:221:8: warning: variable 'sum' set but not used [-Wunused-but-set-variable]
   auto sum = [&](int i, int j){ return SUM[j] - SUM[i]; };
        ^~~
sequence.cpp:114:7: warning: variable 'check_sol' set but not used [-Wunused-but-set-variable]
  auto check_sol = [&](ll x, const vi &t){
       ^~~~~~~~~
sequence.cpp: In instantiation of 'll custom_binary_search(ll, ll, const ll&, Pred, bool) [with Pred = main()::<lambda(ll)>; ll = long long int]':
sequence.cpp:189:60:   required from here
sequence.cpp:36:25: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
    ll mid = low + (high - low >> 1);
                    ~~~~~^~~~~
sequence.cpp:44:25: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
    ll mid = low + (high - low >> 1);
                    ~~~~~^~~~~

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