102672024-03-29 20:57:32111Egyengetőcpp17Accepted 100/1003.375s30060 KiB
#include <bits/stdc++.h>
using namespace std;

#include "grader.h"

#define int long long

struct BST 
{
	struct Node
	{
		int key;
		int count;
		int sum;
		int min;
		int max;
		int size;
		int height;
		int left;
		int right;
		
		Node()
		{
			key = 0;
			count = 0;
			sum = 0;
			size = 0;
			height = 0;
			left = 0;
			right = 0;
		}
		
		Node(int k)
		{
			key = k;
			count = 1;
			sum = k;
			size = 1;
			height = 1;
			left = 0;
			right = 0;
		}
	};
	
	Node* nodes;
	int root;
	int size;
	
	BST(int n)
	{
		nodes = (Node*)malloc((1 + n) * sizeof(Node));
		nodes[0] = Node();
		root = 0;
		size = 0;
	}
	
	void update(int i)
	{
		if (i == 0)
		{
			return;
		}
		nodes[i].sum = nodes[i].key * nodes[i].count + nodes[nodes[i].left].sum + nodes[nodes[i].right].sum;
		nodes[i].size = nodes[i].count + nodes[nodes[i].left].size + nodes[nodes[i].right].size;
		nodes[i].height = 1 + max(nodes[nodes[i].left].height, nodes[nodes[i].right].height);
	}
	
	void rotate_left(int& i)
	{
		int t = nodes[i].right;
		nodes[i].right = nodes[t].left;
		nodes[t].left = i;
		i = t;
	}
	
	void rotate_right(int& i)
	{
		int t = nodes[i].left;
		nodes[i].left = nodes[t].right;
		nodes[t].right = i;
		i = t;
	}
	
	void insert(int& i, int k)
	{
		if (i == 0)
		{
			i = ++size;
			nodes[i] = Node(k);
			return;
		}
		if (k == nodes[i].key)
		{
			nodes[i].count++;
		}
		else
		{
			if (k < nodes[i].key)
			{
				insert(nodes[i].left, k);
			}
			if (k > nodes[i].key)
			{
				insert(nodes[i].right, k);
			}
			int hl = nodes[nodes[i].left].height;
			int hr = nodes[nodes[i].right].height;
			if (hr - hl > 1)
			{
				rotate_left(i);
				update(nodes[i].left);
			}
			if (hl - hr > 1)
			{
				rotate_right(i);
				update(nodes[i].right);
			}
		}
		update(i);
	}
	
	void insert(int k)
	{
		insert(root, k);
	}
	
	int rank_of_key(int i, int k)
	{
		int result = 0;
		while (i != 0) 
		{
			if (nodes[i].key <= k) 
			{
				result += nodes[i].count + nodes[nodes[i].left].size;
				i = nodes[i].right;
			}
			else 
			{
				i = nodes[i].left;
			}
		}
		return result;
	}
	
	int rank_of_key(int k)
	{
		return rank_of_key(root, k);
	}
	
	int find_by_rank(int i, int r)
	{
		while (i != 0)
		{
			if (r <= nodes[nodes[i].left].size)
			{
				i = nodes[i].left;
				continue;
			}
			r -= nodes[nodes[i].left].size + nodes[i].count;
			if (r > 0)
			{
				i = nodes[i].right;
				continue;
			}
			break;
		}
		return i;
	}
	
	int find_by_rank(int r)
	{
		return nodes[find_by_rank(root, r)].key;
	}
	
	int prefix_sum(int i, int k)
	{
		int result = 0;
		while (i != 0) 
		{
			if (nodes[i].key <= k) 
			{
				result += nodes[i].sum - nodes[nodes[i].right].sum;
				i = nodes[i].right;
			}
			else 
			{
				i = nodes[i].left;
			}
		}
		return result;
	}
	
	int prefix_sum(int k)
	{
		return prefix_sum(root, k);
	}
};

signed main(){
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int N=getN(),K=getK(),M=4e8,S=0,C=0;
	BST t(N);
	auto calc=[&](int x)->int{
		int s1=t.prefix_sum(x);
		int c1=t.rank_of_key(x);
		int s2=S-t.prefix_sum(x+K);
		int c2=C-t.rank_of_key(x+K);
		return x*c1-s1+s2-(x+K)*c2;
	};
	for(int i=0;i<N;i++){
		int y=Data();
		t.insert(y);
		S+=y;
		C++;
		int ans=1e18;
		{
			int l=0,h=i;
			while(l<h){
				int m=(l+h)/2;
				if(m>i-t.rank_of_key(t.find_by_rank(m)+K)){
					h=m;
				}
				else{
					l=m+1;
				}
			}
			ans=min(ans,calc(t.find_by_rank(h)));
			ans=min(ans,calc(t.find_by_rank(h+1)));
		}
		{
			int l=0,h=i;
			while(l<h){
				int m=(l+h)/2;
				if(t.rank_of_key(t.find_by_rank(m)-K)>i-m){
					h=m;
				}
				else{
					l=m+1;
				}
			}
			ans=min(ans,calc(t.find_by_rank(h)-K));
			ans=min(ans,calc(t.find_by_rank(h+1)-K));
		}
		Solution(ans);
	}
	return 0;
}
SubtaskSumTestVerdictTimeMemory
subtask10/0
1Accepted3ms1896 KiB
2Accepted3ms2124 KiB
subtask220/20
3Accepted4ms2340 KiB
4Accepted4ms2428 KiB
5Accepted7ms2440 KiB
6Accepted8ms2568 KiB
7Accepted9ms2776 KiB
8Accepted1.368s4444 KiB
subtask320/20
9Accepted14ms3196 KiB
10Accepted14ms3308 KiB
11Accepted14ms3568 KiB
12Accepted17ms3640 KiB
13Accepted16ms3860 KiB
14Accepted14ms4032 KiB
subtask420/20
15Accepted68ms4232 KiB
16Accepted133ms5512 KiB
17Accepted120ms4312 KiB
18Accepted130ms5540 KiB
19Accepted138ms5128 KiB
20Accepted163ms5596 KiB
subtask540/40
21Accepted2.75s22488 KiB
22Accepted1.623s4300 KiB
23Accepted1.85s4276 KiB
24Accepted2.111s4544 KiB
25Accepted3.375s30060 KiB
26Accepted3.355s29932 KiB