Cod sursa(job #2950774)

Utilizator GabiTulbaGabi Tulba-Lecu GabiTulba Data 4 decembrie 2022 17:24:55
Problema Heavy Path Decomposition Scor 80
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 6.07 kb
#include <bits/stdc++.h>
#include <climits>

#define MAX_LOG 2
#define MAX_N 100000

using namespace std;

class SegmentTree {
private:
  vector<int> tree;
  int S;
  int E;
  int Val;
  int Pos;
  int N;

  void buildInternal(int node, int start, int end, vector<int> &array) {
    if(start == end) {
      tree[node] = array[start];
      return;
    }

    int middle = (start + end) / 2;
    buildInternal(2 * node + 1, start, middle, array);
    buildInternal(2 * node + 2, middle + 1, end, array);

    tree[node] = max(tree[2 * node + 1], tree[2 * node + 2]);
  }

  void updateInternal(int node, int start, int end) {
    if(start == end) {
      tree[node] = Val;
      return;
    }

    int middle = (start + end) / 2;
    if(Pos <= middle) {
      updateInternal(2 * node + 1, start, middle);
    } else {
      updateInternal(2 * node + 2, middle + 1, end);
    }

    tree[node] = max(tree[2 * node + 1], tree[2 * node + 2]);
  }

  void queryInternal(int node, int start, int end) {
    if(S <= start && end <= E) {
      Val = max(Val, tree[node]);
      return;
    }

    int middle = (start + end) / 2;
    if(S <= middle) {
      queryInternal(2 * node + 1, start, middle);
    }
    if(E > middle) {
      queryInternal(2 * node + 2, middle + 1, end);
    }
  }

public:
  SegmentTree(vector<int> &array) {
    tree = vector<int>(4 * array.size());
    tree.shrink_to_fit();
    N = array.size();

    buildInternal(0, 0, N - 1, array);
  }

  void update(int pos, int val) {
    Pos = pos;
    Val = val;

    updateInternal(0, 0, N - 1);
  }

  int query(int start, int end) {
    S = start;
    E = end;
    Val = INT_MIN;

    queryInternal(0, 0, N - 1);

    return Val;
  }
};

struct NodeInfo {
  int pathIdx;
  int position;
  int pathParent;
  int father;

  NodeInfo() {
    pathIdx = -1;
    position = -1;
    pathParent = -1;
    father = -1;
  }

  NodeInfo(int pathIdx_, int position_, int pathParent_, int father_) {
    pathIdx = pathIdx_;
    position = position_;
    pathParent = pathParent_;
    father = father_;
  }
};

class HeavyLightDecomp {
private:
  int N;
  vector<SegmentTree> paths;
  vector<NodeInfo> infos;

  void computeSize(int node, vector<int> &size, vector<int> *tree) {
    size[node] = 1;

    for(auto &son: tree[node]) {
      if(son != infos[node].father) {
        infos[son].father = node;
        computeSize(son, size, tree);
        size[node] += size[son];
      }
    }
  }

  int computePaths(int node, vector<int> &size, vector<int> *tree,
                   int *cost, vector<vector<int>> &pathArrays, int pathIdx) {
    int bestSon = -1;

    if((int) pathArrays.size() <= pathIdx) {
      pathArrays.push_back(vector<int>());
    }

    pathArrays[pathIdx].push_back(cost[node]);

    for(auto &son: tree[node]) {
      if(son != infos[node].father && (bestSon == -1 || size[bestSon] < size[son])) {
        bestSon = son;
      }
    }

    if(bestSon == -1) {
      return 0;
    }

    int addedPaths = 0;
    infos[bestSon] = NodeInfo(pathIdx, infos[node].position + 1, infos[node].pathParent, node);
    addedPaths += computePaths(bestSon, size, tree, cost, pathArrays, pathIdx);

    for(auto &son: tree[node]) {
      if(son != infos[node].father && son != bestSon) {
        infos[son] = NodeInfo(pathIdx + addedPaths + 1, 0, son, node);
        addedPaths += computePaths(son, size, tree, cost, pathArrays, pathIdx + addedPaths + 1) + 1;
      }
    }

    return addedPaths;
  }

public:
  HeavyLightDecomp(vector<int> *tree, int *cost, int n) {
    N = n;
    infos = vector<NodeInfo>(N);
    infos.shrink_to_fit();

    vector<int> size = vector<int>(N);
    size.shrink_to_fit();
    vector<vector<int>> pathArrays;
    computeSize(0, size, tree);

    infos[0] = NodeInfo(0, 0, 0, -1);
    computePaths(0, size, tree, cost, pathArrays, 0);
    pathArrays.shrink_to_fit();

    for(auto &pathArray: pathArrays) {
      paths.push_back(SegmentTree(pathArray));
    }
    paths.shrink_to_fit();
  }

  void update(int node, int val) {
    paths[infos[node].pathIdx].update(infos[node].position, val);
  }

  int query(int father, int son) {
    int val = INT_MIN;
    while(infos[father].pathIdx != infos[son].pathIdx) {
      val = max(val, paths[infos[son].pathIdx].query(0, infos[son].position));
      son = infos[infos[son].pathParent].father;
    }
    val = max(val, paths[infos[father].pathIdx].query(infos[father].position, infos[son].position));
    return val;
  }
};

FILE *fin, *fout;
int N, M;
int type, x, y, z;
vector<int> tree[MAX_N];
int cost[MAX_N];
int depth[MAX_N];
int dp[MAX_LOG][MAX_N];

void precalc_lca(int node, int father, int d) {
  dp[0][node] = father;
  depth[node] = d;

  for(int i = 1; i < MAX_LOG; i++) {
    dp[i][node] = dp[i-1][node];
    for(int k = 0; k < 317; k++) {
      dp[i][node] = dp[i-1][dp[i][node]];
    }
  }

  for(auto &son: tree[node]) {
    if(depth[son] == -1) {
      precalc_lca(son, node, d + 1);
    }
  }
}

int lca(int x, int y) {
  if(depth[x] > depth[y]) {
    swap(x, y);
  }
  for(int i = MAX_LOG - 1; i >= 0; i--) {
    for(int k = 0; k < 317; k++) {
      if(depth[dp[i][y]] >= depth[x]) {
        y = dp[i][y];
      }
    }
  }

  if(x == y) {
    return x;
  }

  for (int i = MAX_LOG - 1; i >= 0; i--) {
    for(int k = 0; k < 317; k++) {
      if(dp[i][y] != dp[i][x]) {
        y = dp[i][y];
        x = dp[i][x];
      }
    }
  }

  return dp[0][x];
}

int main() {
  fin = fopen("heavypath.in", "r");
  fout = fopen("heavypath.out", "w");

  fscanf(fin, "%d%d", &N, &M);

  for(int i = 0; i < N; i++) {
    fscanf(fin, "%d", &cost[i]);
  }
  for(int i = 0; i < N -1; i++) {
    fscanf(fin, "%d%d", &x, &y);

    x--;
    y--;
    tree[x].push_back(y);
    tree[y].push_back(x);
  }

  for(int i = 0; i < N; i++) {
    tree[i].shrink_to_fit();
  }

  for(int i = 0; i < N; i++) {
    depth[i] = -1;
  }
  precalc_lca(0, 0, 0);

  HeavyLightDecomp hpd(tree, cost, N);

  for(int i = 0; i < M; i++) {
    fscanf(fin, "%d%d%d", &type, &x, &y);

    if(type == 0) {
      x--;
      hpd.update(x, y);
    } else {
      x--;
      y--;
      z = lca(x, y);

      fprintf(fout, "%d\n", max(hpd.query(z, x), hpd.query(z, y)));
    }
  }

  return 0;
}