10267 2024. 03. 29 20:57:32 111 Egyengető cpp17 Elfogadva 100/100 3.375s 30060 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;
}
Részfeladat Összpont Teszt Verdikt Idő Memória
subtask1 0/0
1 Elfogadva 3ms 1896 KiB
2 Elfogadva 3ms 2124 KiB
subtask2 20/20
3 Elfogadva 4ms 2340 KiB
4 Elfogadva 4ms 2428 KiB
5 Elfogadva 7ms 2440 KiB
6 Elfogadva 8ms 2568 KiB
7 Elfogadva 9ms 2776 KiB
8 Elfogadva 1.368s 4444 KiB
subtask3 20/20
9 Elfogadva 14ms 3196 KiB
10 Elfogadva 14ms 3308 KiB
11 Elfogadva 14ms 3568 KiB
12 Elfogadva 17ms 3640 KiB
13 Elfogadva 16ms 3860 KiB
14 Elfogadva 14ms 4032 KiB
subtask4 20/20
15 Elfogadva 68ms 4232 KiB
16 Elfogadva 133ms 5512 KiB
17 Elfogadva 120ms 4312 KiB
18 Elfogadva 130ms 5540 KiB
19 Elfogadva 138ms 5128 KiB
20 Elfogadva 163ms 5596 KiB
subtask5 40/40
21 Elfogadva 2.75s 22488 KiB
22 Elfogadva 1.623s 4300 KiB
23 Elfogadva 1.85s 4276 KiB
24 Elfogadva 2.111s 4544 KiB
25 Elfogadva 3.375s 30060 KiB
26 Elfogadva 3.355s 29932 KiB