I am getting TLE, in this problem. I think my solution runs in O(2nlogn).
Please help !!!
#include <iostream>
#include <cstdio>
#include <string.h>
#define mod 1000000007
#define S 200100
using namespace std;
int kl;
int kr;
int hleft[S];
int hright[S];
void minHeap(int val)
{
hright[++kr] = val;
int idx = (kr-1)/2;
while(idx>=0){
if(2*idx+1 <=kr){
if(hright[idx] > hright[2*idx+1]){
hright[idx] = hright[idx]^hright[2*idx+1];
hright[2*idx+1] = hright[2*idx+1]^hright[idx];
hright[idx] = hright[idx]^hright[2*idx+1];
}
}
if(2*idx+2 <= kr){
if(hright[idx] > hright[2*idx+2]){
hright[idx] = hright[idx]^hright[2*idx+2];
hright[2*idx+2] = hright[2*idx+2]^hright[idx];
hright[idx] = hright[idx]^hright[2*idx+2];
}
}
if(idx == 0)
return;
idx = (idx-1)>>1;
}
}
void maxHeap(int val)
{
hleft[++kl] = val;
int idx = (kl-1)/2 ;
while(idx >= 0){
if(2*idx+1 <= kl){
if(hleft[idx] < hleft[2*idx+1])
swap(hleft[idx] , hleft[2*idx+1]);
}
if(2*idx+2 <= kl){
if(hleft[idx] < hleft[2*idx+2])
swap(hleft[idx] , hleft[2*idx+2]);
}
if(idx == 0)
return;
idx = (idx-1)>>1;
}
}
int main()
{
int t;
int a,b,c,n;
scanf("%d",&t);
while(t--){
scanf("%d%d%d%d",&a,&b,&c,&n);
memset(hleft,0,sizeof(hleft));
memset(hright,0,sizeof(hright));
long long int sum = 1;
long int val = 1;
hleft[0] = 1;
val = (((a*hleft[0])%mod+(b*2)%mod+c%mod))%mod;
if(val < 0)
val += mod;
sum += val;
hright[0] = val;
kl = 0;
kr = 0;
for(int i=2;i<n;i++){
val = (((a*hleft[0])%mod + (b*(i+1))%mod + c%mod))%mod;
if(val < 0)
val += mod;
sum += val;
minHeap(val);
if(!(i&1)){
maxHeap(hright[0]);
hright[0] = 1000000107;
}
else{
if(hright[0] <= hleft[0]){
hleft[0] = hleft[0]^hright[0];
hright[0] = hright[0]^hleft[0];
hleft[0] = hleft[0]^hright[0];
maxHeap(-1);
}
}
}
printf("%lld\n",sum);
}
return 0;
}
I have used 2 heaps. first for left elements to median(root) in descending order and second after median to right elements in ascending order.
Thanx in advance