HackerEarth Treat in a restaurant problem solution

In this HackerEarth Treat in a restaurant problem solution Your city has N restaurants that are located in the form of a tree and connected by roads. This tree is a connected undirected graph consisting of N vertices that are connected by N - 1 edges. Each restaurants charges a fixed amount irrespective of the amount of food that you eat. You are given the amount A[i] for every 1 <= i <= N. 

You are giving a treat to your friends at four different restaurants with no limit on the amount of food that they can eat. 

You start your journey from the restaurant located at the root to the restaurant that is located at any of the leaves through a simple path. Your friends can select any four different restaurants P, Q, R, and S that are situated on the simple path. Here, P is the ancestor of Q, Q is the ancestor of R, and R is the ancestor of S. Your friends wants you to spend as much money as possible. The sum of money A[P] + A[Q] + A[R] + A[S] must be the maximum amount. The amount that you spend in four restaurants must be strictly increasing such as A[P] < A[Q] < A[R] < A[S].

Print the maximum total amount of money that you can spend by following the condition. Otherwise, print 0.

HackerEarth Treat in a restaurant problem solution.

#include "bits/stdc++.h"
using namespace std;
 
#define int long long
#define F first
#define S second
#define siz(x) ((int)x.size())
#define rep(i,a,n) for (int i = a; i <= n; ++i)
#define all(v)  v.begin(), v.end()
#define pb push_back
#define P pair < int, int >
#define E cout << '\n'

const int mod = 1e9 + 7;
const int N = 100000 + 5;

int a[N], sz[N];
int dp[N], idx[N], ctr;
int arr[N], n;
multiset<int>st;
vector<int>v[N];

class merge_sort_tree{
    public:
    vector < P > tree[N << 2];
    vector<int>suff_max[N<<2];    
    void build(int l, int r, int node){
        if (l > r)return;
        if(l == r){
            int x = arr[l];
            tree[node].pb({a[x], x});
            return ;
        }
        int mid = l + r >> 1, lc = node + node, rc = 1 + lc;
        build(l, mid, lc);    build(mid + 1, r, rc);
        merge(all(tree[lc]), all(tree[rc]), back_inserter(tree[node]));
    }
    int query(int l, int r, int ql, int qr, int val, int node){
        if(qr < l || r < ql || l > r || ql > qr)
            return 0;
        if(ql <= l and r <= qr){           
            if (tree[node].empty() || tree[node].back().first <= val) return 0;
            int lo = 0, hi = siz(tree[node])-1, mid;
            while (lo <= hi) {
                mid = (lo+hi)/2;
                if (tree[node][mid].first > val) hi = mid - 1;
                else lo = mid+1;
            }
            assert(hi+1>=0 && hi+1<siz(tree[node]));
           return suff_max[node][hi+1];
        }
        int mid = l + r >> 1, lc = node + node, rc = 1 + lc;

        return max(query(l, mid, ql, qr, val, lc),query(mid + 1, r, ql, qr, val, rc));
    }
    void make_suffix_max() {
        for (int i = 1; i < 4*N; ++i) {
            if (tree[i].empty()) continue;
            suff_max[i].resize(siz(tree[i]));
            suff_max[i].back() = dp[tree[i].back().second];
            for (int j = siz(tree[i])-2; j >= 0; --j) {
              suff_max[i][j] = max(dp[tree[i][j].second], suff_max[i][j+1]);
            }
        }
    }
};

merge_sort_tree obj; 

int dfs(int node, int p = 0) {
    idx[node] = ++ctr;
    arr[ctr] = node;
    sz[node] = 1;   
    int mx = a[node];
    for (int x: v[node]) {
        if (x != p) {
            mx = max(mx, dfs(x, node)); 
            sz[node] += sz[x];
        }
    }
    return dp[node] = mx;
}

int dfs2(int x, int p) {
    st.insert(a[x]);
    int ans = 0;
    if (dp[x]) {
       auto it = st.lower_bound(a[x]);
       if (it != st.begin()) {
        int ok = obj.query(1, n, idx[x]+1, idx[x]+sz[x]-1, a[x], 1);
        if (ok) {
            --it;
            ans = max(ans, ok+a[x]+*it);
        }
       }
    }
    for (int ok: v[x]) {
        if (ok != p) {
            ans = max(ans, dfs2(ok,x));
        }
    }
    st.erase(st.find(a[x]));
    return ans;
}

inline void solve() {
   int ans = 0;
   cin >> n;
   rep(i,1,n) {
    cin >> a[i];
   }
   int l, r;
   rep(i,2,n) {
    cin >> l >> r;
    v[l].push_back(r);
    v[r].push_back(l);
   }   
   dfs(1);
   rep(i,1,n) {
    dp[i] += (dp[i]==a[i] ? -dp[i]:a[i]);
   }   
   obj.build(1,n,1);
   obj.make_suffix_max();
   cout << dfs2(1, 0);  
}
signed main() {
  ios_base::sync_with_stdio(0);
  cin.tie(NULL);
  cout.tie(NULL);
  int t = 1;
  //cin >> t; while(t--)
  solve();
  return 0;
}