BST maintenance

                            Solution in C :

In   C++  :






#include <bits/stdc++.h>
using namespace std;

const int
MAXV = 250005;

int V;
int p[MAXV];
int pos[MAXV];

set< int > S;
int L[MAXV];
int R[MAXV];
int P[MAXV];

bool seen[MAXV];
vector< int > childs;
int depth[MAXV];
int size[MAXV];
vector< int > G[MAXV];
int root;

long long sum[MAXV];

long long nodeCount[3];
long long distancesSum[3];

int branch[MAXV];
int curBranch;

void computeSizes(int u)
{
seen[u] = 1;
size[u] = 1;

for (int v : G[u]) if (!seen[v])
{
computeSizes(v);
size[u] += size[v];
}

seen[u] = 0;
}

void dfs(int u)
{
branch[u] = curBranch;
seen[u] = 1;
childs.push_back(u);

for (int v : G[u]) if (!seen[v])
{
depth[v] = depth[u] + 1;
dfs(v);
}

seen[u] = 0;
}

int findCenter(int u)
{
bool isCenter = 
((size[root] - size[u]) * 2 <= size[root]);

for (int v : G[u]) if (!seen[v])
isCenter &= (size[v] * 2 <= size[root]);

if (isCenter)
return u;

seen[u] = 1;

for (int v : G[u]) if (!seen[v])
{
int c = findCenter(v);
if (c != -1)
{
seen[u] = 0;
return c;
}
}

seen[u] = 0;

return -1;
}

bool posCmp(int a, int b)
{
return pos[a] < pos[b];
}

void solve(int u)
{
computeSizes(u);

if (size[u] == 1) // single leaf
return ;

root = u;
u = findCenter(u);

seen[u] = 1;

childs.clear();
curBranch = 0;
for (int v : G[u]) if (!seen[v]) {

depth[v] = 1;
dfs(v);

// fill 0
nodeCount[curBranch] = 0;
distancesSum[curBranch] = 0;

curBranch++;
}

assert(curBranch <= 3);

childs.push_back(u); // root

// sort can be avoided
sort(childs.begin(), childs.end(), posCmp);

bool isRootAdded = false;

for (int c : childs)
{
if (c == u) // add root
{
for (int i = 0; i < curBranch; ++i)
sum[ pos[c] ] += distancesSum[i];

isRootAdded = true;
}
else {
int b = branch[c];
if (isRootAdded)
{
for (int i = 0; i < curBranch; ++i) if (i != b)
    sum[ pos[c] ] += 
    distancesSum[i] + nodeCount[i] * depth[c];
// add root
sum[ pos[c] ] += depth[c];
}
distancesSum[b] += depth[c];
nodeCount[b]++;
}
}

for (int v : G[u]) if (!seen[v])
solve(v);

seen[u] = 0;
}

int main()
{
ios_base::sync_with_stdio(0);
cin.tie(0);

cin >> V;
memset(L, -1, sizeof(L));
memset(R, -1, sizeof(R));

for (int i = 0; i < V; ++i)
{
cin >> p[i];
--p[i];

pos[ p[i] ] = i;

auto lb = S.lower_bound(p[i]);
if (lb != S.end())
{
if (L[*lb] < 0) {
L[*lb] = p[i];
G[*lb].push_back(p[i]);
G[p[i]].push_back(*lb);
}
}

if (lb != S.begin()) {
--lb;
if (R[*lb] < 0) {
R[*lb] = p[i];
G[*lb].push_back(p[i]);
G[p[i]].push_back(*lb);
}    
}
S.insert(p[i]);
}

solve(p[0]);

for (int i = 1; i < V; ++i)
sum[i] += sum[i - 1];

for (int i = 0; i < V; ++i)
cout << sum[i] << '\n';

return 0;
}









In   Java :






import java.io.*;
import java.util.*;

public class Solution {

BufferedReader br;
PrintWriter out;
StringTokenizer st;
boolean eof;

static int getSize(Node v) {
return v == null ? 0 : v.size;
}

static class Node implements Comparable<Node>
 {
Node left, right;
int val;
int depth;

int size;

int pathNum, inPathPos;

@Override
public int compareTo(Node o) {
return Integer.compare(val, o.val);
}

public Node(int val) {
this.val = val;
}

@Override
public String toString() {
StringBuilder sb = new StringBuilder();
if (left != null) {
sb.append('[');
sb.append(left.toString());
sb.append("] ");
}
sb.append(val);
if (right != null) {
sb.append(" [");
sb.append(right.toString());
sb.append(']');
}
return sb.toString();
}
}

int nPaths = 0;
List<Node> nextNode;
int[] pathSize;
FenwickTree[] fen;

void addUptoRoot(Node v) {
do {
int pathNum = v.pathNum;
int inPathPos = v.inPathPos;
fen[pathNum].add(0, inPathPos);
v = nextNode.get(pathNum);
} while (v != null);
}

long getUptoRoot(Node v) {
long ret = 0;
do {
int pathNum = v.pathNum;
int inPathPos = v.inPathPos;
ret += fen[pathNum].get(0, inPathPos);
v = nextNode.get(pathNum);
} while (v != null);
return ret;
}

void solve() throws IOException {
int n = nextInt();
TreeSet<Node> set = new TreeSet<>();
Node root = null;
Node[] order = new Node[n];

for (int i = 0; i < n; i++) {
int x = nextInt();
Node v = new Node(x);
order[i] = v;
set.add(v);
if (set.size() == 1) {
root = v;
continue;
}
Node before = set.lower(v);
Node after = set.higher(v);
if (before != null
&& before.right == null
&& (after == null || 
after.left != null || before.depth < after.depth))
 {
before.right = v;
v.depth = before.depth + 1;
} else {
after.left = v;
v.depth = after.depth + 1;
}
}

dfs1(root);

nextNode = new ArrayList<Node>(n);
pathSize = new int[n];
dfs2(root, 0, 0, null);



fen = new FenwickTree[nPaths];
for (int i = 0; i < nPaths; i++) {
fen[i] = new FenwickTree(pathSize[i]);
}

long outp = 0;
long sumDistRoot = 0;

for (int i = 0; i < n; i++) {
Node v = order[i];
sumDistRoot += v.depth;
addUptoRoot(v);
long sum = getUptoRoot(v) - (i + 1); 
long delta = (long) (i + 1) * v.depth - 2 * sum;
outp += sumDistRoot + delta;
out.println(outp);
}

}

void dfs1(Node v) {
if (v == null) {
return;
}
dfs1(v.left);
dfs1(v.right);
v.size = getSize(v.left) + getSize(v.right) + 1;
}

void dfs2(Node v, int pathNum, int inPathPos, Node par) {
if (v == null) {
return;
}
if (inPathPos == 0) {
pathNum = nPaths++;
nextNode.add(par);
}
v.pathNum = pathNum;
v.inPathPos = inPathPos;
if (v.left == null && v.right == null) {
pathSize[pathNum] = inPathPos + 1;
return;
}
if (getSize(v.left) > getSize(v.right)) {
dfs2(v.left, pathNum, inPathPos + 1, v);
dfs2(v.right, -1, 0, v);
} else {
dfs2(v.right, pathNum, inPathPos + 1, v);
dfs2(v.left, -1, 0, v);
}
}

static class IntList {
/**
* Never shrinks
*/

private int[] data;
private int size;

public IntList(int cap) {
data = new int[cap];
size = 0;
}

void add(int x) {
data[size++] = x;
}

void clear() {
size = 0;
}

int get(int idx) {
if (idx < 0 || idx >= size) {
throw new IndexOutOfBoundsException();
}
return data[idx];
}
}

static class FenwickTree {
private int n;

private long[] c0;
private long[] c1;

public FenwickTree(int n) {
this.n = n;
this.c0 = new long[n];
this.c1 = new long[n];
}

void add(int low, int high) {
/**
* [low, high]
*/
internalUpdate(low, -(low - 1), 1);
internalUpdate(high, high, -1);
}

private void internalUpdate(int x, int d0, int d1) {
for (int i = x; i < n; i |= i + 1) {
c0[i] += d0;
c1[i] += d1;
}
}

long get(int low, int high) {
/**
* [low, high]
*/
return get(high) - get(low - 1);
}

long get(int x) {
/**
* [0, x]
*/
long a1 = 0;
long a0 = 0;
for (int i = x; i >= 0; i = (i & (i + 1)) - 1) {
a1 += c1[i];
a0 += c0[i];
}
return a1 * x + a0;
}
}

Solution() throws IOException {
br = new BufferedReader(new InputStreamReader(System.in));
out = new PrintWriter(System.out);
solve();
out.close();
}

public static void main(String[] args) throws IOException {
new Solution();
}

String nextToken() {
while (st == null || !st.hasMoreTokens()) {
try {
st = new StringTokenizer(br.readLine());
} catch (Exception e) {
eof = true;
return null;
}
}
return st.nextToken();
}

String nextString() {
try {
return br.readLine();
} catch (IOException e) {
eof = true;
return null;
}
}

int nextInt() throws IOException {
return Integer.parseInt(nextToken());
}

long nextLong() throws IOException {
return Long.parseLong(nextToken());
}

double nextDouble() throws IOException {
return Double.parseDouble(nextToken());
}
}









In    C  :







#include <stdio.h>
#include <stdlib.h>
typedef struct _node{
int x;
struct _node *next;
} lnode;
void init( int n ,int *tree);
void range_increment( int i, int j,
 int val ,int *tree);
int query( int i ,int *tree);
void insert_edge(int x,int y);
void dfs0(int u);
void preprocess();
int lca(int a,int b);
int dist(int u,int v);
void dfs1(int u,int p);
int dfs2(int u,int p);
void decompose(int root,int p);
int a[250000],cut[250000]={0},parent[250000],
DP[18][250000],mid[750000],left[750000],
right[750000],level[250000],
sub[250000],N,NN,nn;
long long count[250000]={0},
sum[250000]={0},con[250000]={0};
lnode *table[250000]={0};

int main(){
int x,y,z,leftd,rightd,i;
long long ans,aa=0;
scanf("%d",&NN);
for(i=0;i<NN;i++)
scanf("%d",a+i);
init(NN,mid);
init(NN,left);
init(NN,right);
for(i=0;i<NN;i++){
leftd=x=query(a[i]-1,left);
if(!x)
leftd=1;
rightd=y=query(a[i]-1,right);
if(!y)
rightd=NN;
z=query(a[i]-1,mid);
if(z)
insert_edge(z-1,a[i]-1);
range_increment(leftd-1,rightd-1,a[i]-z,mid);
range_increment(a[i]-1,rightd-1,a[i]-x,left);
range_increment(leftd-1,a[i]-1,a[i]-y,right);
}
preprocess();
decompose(a[NN/2]-1,-1);
for(i=0;i<NN;i++){
for(ans=sum[a[i]-1],x=a[i]-1;1;x=parent[x]){
if(parent[x]==-1)
break;
ans+=sum[parent[x]]-con[x]+dist(a[i]-1,
parent[x])*(count[parent[x]]-count[x]);
}
for(x=a[i]-1;x!=-1;x=parent[x]){
sum[x]+=dist(a[i]-1,x);
count[x]++;
if(parent[x]!=-1)
con[x]+=dist(a[i]-1,parent[x]);
}
printf("%lld\n",aa+=ans);
}
return 0;
}
void init( int n ,int *tree){
N = 1;
while( N < n ) N *= 2;
int i;
for( i = 1; i < N + n; i++ ) tree[i] = 0;
}
void range_increment( int i, int j, 
int val ,int *tree){
for( i += N, j += N; i <= j; 
i = ( i + 1 ) / 2, j = ( j - 1 ) / 2 )
{
if( i % 2 == 1 ) tree[i] += val;
if( j % 2 == 0 ) tree[j] += val;
}
}
int query( int i ,int *tree){
int ans = 0,j;
for( j = i + N; j; j /= 2 ) ans += tree[j];
return ans;
}
void insert_edge(int x,int y){
lnode *t=malloc(sizeof(lnode));
t->x=y;
t->next=table[x];
table[x]=t;
t=malloc(sizeof(lnode));
t->x=x;
t->next=table[y];
table[y]=t;
return;
}
void dfs0(int u){
lnode *x;
for(x=table[u];x;x=x->next)
if(x->x!=DP[0][u]){
DP[0][x->x]=u;
level[x->x]=level[u]+1;
dfs0(x->x);
}
return;
}
void preprocess(){
int i,j;
level[a[0]-1]=0;
DP[0][a[0]-1]=a[0]-1;
dfs0(a[0]-1);
for(i=1;i<18;i++)
for(j=0;j<NN;j++)
DP[i][j] = DP[i-1][DP[i-1][j]];
return;
}
int lca(int a,int b){
int i;
if(level[a]>level[b]){
i=a;
a=b;
b=i;
}
int d = level[b]-level[a];
for(i=0;i<18;i++)
if(d&(1<<i))
b=DP[i][b];
if(a==b)return a;
for(i=17;i>=0;i--)
if(DP[i][a]!=DP[i][b])
a=DP[i][a],b=DP[i][b];
return DP[0][a];
}
int dist(int u,int v){
return level[u] + level[v] - 2*level[lca(u,v)];
}
void dfs1(int u,int p){
sub[u]=1;
nn++;
lnode *x;
for(x=table[u];x;x=x->next)
if(x->x!=p && !cut[x->x]){
dfs1(x->x,u);
sub[u]+=sub[x->x];
}
return;
}
int dfs2(int u,int p){
lnode *x;
for(x=table[u];x;x=x->next)
if(x->x!=p && sub[x->x]>nn/2 && !cut[x->x])
return dfs2(x->x,u);
return u;
}
void decompose(int root,int p){
nn=0;
dfs1(root,root);
int centroid = dfs2(root,root);
parent[centroid]=p;
cut[centroid]=1;
lnode *x;
for(x=table[centroid];x;x=x->next)
if(!cut[x->x])
decompose(x->x,centroid);
return;
}