HackerEarth Mathison and the peculiar sums problem solution

In this HackerEarth Mathison and the problem of the peculiar sum solution, Mathison has bought a new tree from the local market. There was a sale at the market and the tree came with a free set of weights, one weight for each one of its nodes. Mathison is getting bored of just looking at his tree (even though the tree is quite lovely) and starts doing some computations using the weights and the tree's structure.

He starts by choosing a node and getting all the weights on the path from that node to the root (i.e. the node with label 1). He then sorts the weights in increasing order. Next, he sums all the weights that appear in even positions (in the sorted sequence of weights) and then all the weights that appear in odd positions (in the same sorted sequence). The peculiar value of the node is defined to be the product of these two sums.

Mathison performs this calculation for all nodes and wants you to help him make sure that it's correct. Because the numbers are quite large, every peculiar value is computed mod 10^9 + 7.

HackerEarth Mathison and the peculiar sums problem solution.

#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>
#include <map>
#include <unordered_map>

using namespace std;

constexpr int MAX_N = 100'000 + 1;
constexpr int MOD = 1'000'000'000 + 7;

template <class T>
struct OrderedNormalization{
private:
  map<T, int> norm;
public:

  OrderedNormalization(std::vector<T> v){
    sort(v.begin(), v.end());

    int value = 0;
    for (auto x: v)
      norm[x] = ++value;
  }

  int get(T key) const{
    auto it = norm.find(key);
    assert (it != norm.end());
    return it->second;
  }
};

struct Node{
  int sum_evens, sum_odds;
  int cnt;
};

int add_mod(int x, int y){
  int sol = x + y;

  if (sol >= MOD)
    sol -= MOD;

  return sol;
}

int mult_mod(int x, int y){
  return (1LL * x * y) % MOD;
}

Node merge_nodes(Node L, Node R){
  Node answer;
  answer.cnt = L.cnt + R.cnt;

  if (L.cnt % 2 == 0){
    answer.sum_evens = add_mod(L.sum_evens, R.sum_evens);
    answer.sum_odds = add_mod(L.sum_odds, R.sum_odds);
  }
  else{
    answer.sum_evens = add_mod(L.sum_evens, R.sum_odds);
    answer.sum_odds = add_mod(L.sum_odds, R.sum_evens);
  }

  return answer;
}

Node seg_tree[4 * MAX_N];
int key_node[MAX_N];
int pos_node[MAX_N];
int answer_node[MAX_N];
std::vector<int> tree[MAX_N];

void update(int node, int l, int r, int pos, int key, bool add){
  if (l == r){
    if (add){
      seg_tree[node].cnt = 1;
      seg_tree[node].sum_evens = key;
      seg_tree[node].sum_odds = 0;
    }
    else{
      seg_tree[node].cnt = 0;
      seg_tree[node].sum_evens = 0;
      seg_tree[node].sum_odds = 0;
    }
  }
  else{
    int m = (l + r) / 2;

    if (pos <= m)
      update(2 * node, l, m, pos, key, add);
    else
      update(2 * node + 1, m + 1, r, pos, key, add);

    seg_tree[node] = merge_nodes(seg_tree[2 * node], seg_tree[2 * node + 1]);
  }
}

void dfs(int node, int father, const int N){
  update(1, 1, N, pos_node[node], key_node[node], true); // add key
  answer_node[node] = mult_mod(seg_tree[1].sum_evens, seg_tree[1].sum_odds);

  for (int v: tree[node])
    if (v != father)
      dfs(v, node, N);

  update(1, 1, N, pos_node[node], key_node[node], false); // remove key
}

int main() {
  ios_base::sync_with_stdio(false);

  int N;
  cin >> N;

  std::vector<pair<int,int>> keys;
  for (int i = 1; i <= N; i++){
    cin >> key_node[i];
    keys.push_back({key_node[i], i});
  }

  for (int i = 1; i < N; i++){
    int u, v;
    cin >> u >> v;

    tree[u].push_back(v);
    tree[v].push_back(u);
  }

  OrderedNormalization<pair<int,int>> ord_norm(keys);
  for (int i = 1; i <= N; i++)
    pos_node[i] = ord_norm.get({key_node[i], i});

  dfs(1, 0, N);

  for (int i = 1; i <= N; i++)
    cout << answer_node[i] << "\n";

  return 0;
}

Second solution

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

#define L(i) 2 * i
#define R(i) 2 * i + 1
#define MOD 1000000007
#define N 100005
#define rank _rank

typedef struct ${
    int cnt, sum_o, sum_e;
} Node;

pair<int, int> M[N];
int n, curr, val[N], rank[N], ans[N];
vector<int> v[N];
Node tree[4 * N];

void update(int node, int l, int r, int idx, int adder){
    if(l == r){
        tree[node].cnt = adder;
        tree[node].sum_e = adder * val[idx];
        tree[node].sum_o = 0;
        return;
    }
    int m = (l + r) >> 1;
    (rank[idx] <= m) ? update(L(node), l, m, idx, adder) : update(R(node), m + 1, r, idx, adder);

    tree[node].cnt = tree[L(node)].cnt + tree[R(node)].cnt;
    if(tree[L(node)].cnt & 1){
        tree[node].sum_e = (tree[L(node)].sum_e + tree[R(node)].sum_o) % MOD;
        tree[node].sum_o = (tree[L(node)].sum_o + tree[R(node)].sum_e) % MOD;
    }
    else{
        tree[node].sum_e = (tree[L(node)].sum_e + tree[R(node)].sum_e) % MOD;
        tree[node].sum_o = (tree[L(node)].sum_o + tree[R(node)].sum_o) % MOD;
    }
    return;
}

void dfs(int start, int parent = 0){
    update(1, 1, n, start, 1);
    ans[start] = (1ll * tree[1].sum_o * tree[1].sum_e) % MOD;
    for(auto next : v[start])
        if(next != parent)
            dfs(next, start);
    update(1, 1, n, start, 0);
}

int main(){
    int i,j,a,b;
    scanf("%d", &n);
    for(i = 0; i < n; i++){
        scanf("%d", &val[i + 1]);
        M[i] = {val[i + 1], i + 1};
    }
    sort(M, M + n);
    for(i = 0; i < n; i++)
        rank[M[i].second] = ++curr;

    for(i = 1; i < n; i++){
        scanf("%d %d", &a, &b);
        v[a].push_back(b);
        v[b].push_back(a);
    }
    dfs(1);
    for(i = 1; i <= n; i++)
        printf("%d\n", ans[i]);
    return 0;
}