Cod sursa(job #2031838)

Utilizator preda.andreiPreda Andrei preda.andrei Data 3 octombrie 2017 21:40:39
Problema Heavy Path Decomposition Scor 40
Compilator cpp Status done
Runda Arhiva educationala Marime 5.29 kb
#include <fstream>
#include <vector>

using namespace std;

class SegmentTree
{
public:
  SegmentTree(size_t size)
    : nodes_(vector<int>(4 * size + 4, 0)),
      size_(size) { }

  int Query(int left, int right) const { return Query(1, 1, size_, left, right); }

  void Update(int pos, int value) { Update(1, 1, size_, pos, value); }

private:
  vector<int> nodes_;
  size_t size_;

  int Query(int node, int left, int right, int x, int y) const;

  void Update(int node, int left, int right, int pos, int value);

  static int LeftSon(int pos) { return pos << 1; }
  static int RightSon(int pos) { return (pos << 1) + 1; }
};

int SegmentTree::Query(int node, int left, int right, int x, int y) const
{
  if (x <= left && right <= y) {
    return nodes_[node];
  }

  auto mid = left + (right - left) / 2;
  auto res = 0;

  if (x <= mid) {
    res = max(res, Query(LeftSon(node), left, mid, x, y));
  }
  if (mid < y) {
    res = max(res, Query(RightSon(node), mid + 1, right, x, y));
  }
  return res;
}

void SegmentTree::Update(int node, int left, int right, int pos, int value)
{
  if (left == right) {
    nodes_[node] = value;
    return;
  }

  auto mid = left + (right - left) / 2;
  auto left_son = LeftSon(node);
  auto right_son = RightSon(node);

  if (pos <= mid) {
    Update(left_son, left, mid, pos, value);
  } else {
    Update(right_son, mid + 1, right, pos, value);
  }

  nodes_[node] = max(nodes_[left_son], nodes_[right_son]);
}

struct Node
{
  int value;
  int time = 0;
  int height = 0;
  int weight = 0;
  int special = -1;
  int path_id = -1;
  int path_pos = -1;

  vector<int> edges;
  vector<int> ancestors;
};

using Tree = vector<Node>;

struct Path
{
  SegmentTree values;
  size_t size;
  int root;

  Path(size_t size, int root)
    : values(SegmentTree(size)),
      size(size),
      root(root) {}
};

vector<int> FindAncestors(const Tree &t, int root, int father)
{
  vector<int> ancestors {father};
  int order = 0;

  while (order < (int)t[father].ancestors.size()) {
    ancestors.push_back(t[father].ancestors[order]);
    father = ancestors.back();
    order += 1;
  }
  return ancestors;
}

void Dfs(Tree &t, int root, int father = -1)
{
  if (father != -1) {
    t[root].ancestors = FindAncestors(t, root, father);
  }

  int special = -1;
  for (const auto &son : t[root].edges) {
    if (son != father) {
      Dfs(t, son, root);
      t[root].weight += t[son].weight;
      t[root].height = max(t[root].height, t[son].height);

      if (special == -1 || t[son].weight > t[special].weight) {
        special = son;
      }
    }
  }

  t[root].height += 1;
  t[root].weight += 1;
  t[root].special = special;

  static int time = 0;
  t[root].time = ++time;
}

void Decompose(Tree &t, int node, int father, bool make_new, vector<Path> &paths)
{
  if (make_new) {
    paths.push_back(Path(t[node].height, father));
    t[node].path_id = paths.size() - 1;
    t[node].path_pos = 1;
  } else {
    t[node].path_id = t[father].path_id;
    t[node].path_pos = t[father].path_pos + 1;
  }

  if (t[node].special == -1) {
    return;
  }

  Decompose(t, t[node].special, node, false, paths);

  for (const auto &son : t[node].edges) {
    if (son != father && son != t[node].special) {
      Decompose(t, son, node, true, paths);
    }
  }
}

inline vector<Path> Decompose(Tree &t, int root)
{
  vector<Path> paths;
  Decompose(t, root, -1, true, paths);
  return paths;
}

void UpdateValue(Tree &t, vector<Path> &p, int node, int value)
{
  auto id = t[node].path_id;
  auto pos = t[node].path_pos;
  p[id].values.Update(pos, value);
  t[node].value = value;
}

int FindEarlyAncestor(const Tree &t, int node, int max_time)
{
  auto ancestor = t[node].ancestors[0];
  for (const auto &anc : t[node].ancestors) {
    if (t[anc].time <= max_time) {
      ancestor = anc;
    }
  }
  return ancestor;
}

int FindLca(const Tree &t, int a, int b)
{
  if (t[a].time > t[b].time) {
    swap(a, b);
  }

  while (t[a].time < t[b].time) {
    a = FindEarlyAncestor(t, a, t[b].time);
  }
  return a;
}

int FindPathMax(const Tree &t, const vector<Path> &p, int anc, int son)
{
  int res = 0;

  while (t[son].path_id != t[anc].path_id) {
    auto id = t[son].path_id;
    auto pos = t[son].path_pos;
    res = max(res, p[id].values.Query(1, pos));
    son = p[id].root;
  }

  auto id = t[son].path_id;
  auto left = t[anc].path_pos;
  auto right = t[son].path_pos;
  res = max(res, p[id].values.Query(left, right));

  return res;
}

inline int FindMax(const Tree &t, const vector<Path> &p, int a, int b)
{
  auto lca = FindLca(t, a, b);
  return max(FindPathMax(t, p, lca, a), FindPathMax(t, p, lca, b));
}

int main()
{
  ifstream fin("heavypath.in");
  ofstream fout("heavypath.out");

  int nodes, queries;
  fin >> nodes >> queries;

  Tree tree(nodes);
  for (auto &node : tree) {
    fin >> node.value;
  }

  for (int i = 1; i < nodes; ++i) {
    int x, y;
    fin >> x >> y;
    tree[x - 1].edges.push_back(y - 1);
    tree[y - 1].edges.push_back(x - 1);
  }

  Dfs(tree, 0);

  auto paths = Decompose(tree, 0);

  for (int i = 0; i < nodes; ++i) {
    UpdateValue(tree, paths, i, tree[i].value);
  }

  for (int i = 0; i < queries; ++i) {
    int type, a, b;
    fin >> type >> a >> b;

    if (type == 0) {
      UpdateValue(tree, paths, a - 1, b);
    } else {
      auto res = FindMax(tree, paths, a - 1, b - 1);
      fout << res << "\n";
    }
  }

  return 0;
}