Saturday, June 22, 2019

[Codechef] TAQTREE - Queries On Tree

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : TAQTREE - Queries On Tree
Source            : Codechef
Category          : Graph Theory, Tree
Algorithm         : Easiest HLD Over Subtree and Paths 
Verdict           : Accepted

#include "bits/stdc++.h"
 
using namespace std;
 
static const int maxn = 2e5 + 5;
static const int logn = 19;
 
struct Edge
{
      int v, w, id;
      Edge(int v = 0, int w = 0, int id = 0) : v(v), w(w), id(id) {}
};
 
vector <Edge> graph[maxn];
vector <int> Tree[maxn];
int subtreeSize[maxn], Id[maxn];
int depth[maxn], father[maxn][logn];
int nxt[maxn], in[maxn], out[maxn], tym, nodeCost[maxn];
 
void makeTree(int u = 1, int p = -1)
{
      for (Edge e : graph[u])
      {
            int v = e.v;
            int w = e.w;
            int id = e.id;
            if (v == p) continue;
            Tree[u].push_back(v);
            nodeCost[v] = w;
            Id[id] = v;
            makeTree(v, u);
      }
}
 
void dfsSize(int u = 1)
{
      for (int i = 1; i < logn; i++) father[u][i] = father[ father[u][i-1] ][i-1];
      for (int v : Tree[u])
      {
            dfsSize(v);
            subtreeSize[u] += subtreeSize[v];
            if (subtreeSize[v] > subtreeSize[ Tree[u][0] ]) swap(v, Tree[u][0]);
      }
}
 
void hld(int u = 1)
{
      in[u] = ++tym;
      for (int v : Tree[u])
      {
            nxt[v] = v == Tree[u][0] ? nxt[u] : v;
            hld(v);
      }
      out[u] = tym;
}
 
void dfs(int u = 1)
{
      for (int i = 1; i < logn; i++) father[u][i] = father[ father[u][i-1] ][i-1];
      for (int v : Tree[u])
      {
            depth[v] = depth[u] + 1;
            father[v][0] = u;
            dfs(v);
      }
}
 
int findLCA(int u, int v)
{
      if (depth[u] < depth[v]) swap(u, v);
      for (int i = logn - 1; i >= 0; i--)
      {
            if (depth[ father[u][i] ] >= depth[v]) u = father[u][i];
      }
      if (u == v) return u;
      for (int i = logn - 1; i >= 0; i--)
      {
            if (father[u][i] != father[v][i])
            {
                  u = father[u][i];
                  v = father[v][i];
            }
      }
      return father[u][0];
}
 
int segTree[maxn * 4];
int arr[maxn], N;
 
void update(int node, int a, int b, int pos, int val)
{
      if (a > pos or b < pos) return;
      if (a >= pos and b <= pos)
      {
            segTree[node] = val;
            return;
      }
      int lft = node << 1;
      int rgt = lft | 1;
      int mid = a + b >> 1;
      update(lft, a, mid, pos, val);
      update(rgt, mid + 1, b, pos, val);
      segTree[node] = segTree[lft] + segTree[rgt];
}
 
int query(int node, int a, int b, int i, int j)
{
      if (a > j or b < i) return 0;
      if (a >= i and b <= j) return segTree[node];
      int lft = node << 1;
      int rgt = lft | 1;
      int mid = a + b >> 1;
      int p = query(lft, a, mid, i, j);
      int q = query(rgt, mid + 1, b, i, j);
      return p + q;
}
 
int queryUp(int u, int v)
{
      // u = niche, v = lca
      int sum = 0;
      int L = in[v];
      while (true)
      {
            int l = in[ nxt[u] ];
            int r = in[u];
            if (l <= L)
            {
                  l = max(l, L);
                  sum += query(1, 1, N, l, r);
                  break;
            }
            sum += query(1, 1, N, l, r);
            u = father[ nxt[u] ][0];
      }
      return sum;
}
 
int queryhld(int u, int v)
{
      int lca = findLCA(u, v);
      int lft = queryUp(u, lca);
      int rgt = queryUp(v, lca);
      int middle = queryUp(lca, lca);
      return queryUp(u, lca) + queryUp(v, lca) - 2 * queryUp(lca, lca);
}
 
void updatehld(int u, int val)
{
      nodeCost[u] = val;
      update(1, 1, N, in[u], val);
}
 
int main()
{
      ios_base::sync_with_stdio(false);
      cin.tie(nullptr);
      cout.tie(nullptr);
 
      #ifndef ONLINE_JUDGE
            freopen("in.txt", "r", stdin);
      #endif // ONLINE_JUDGE
 
 
      int tnode;
      cin >> tnode;
      for (int e = 1; e < tnode; e++)
      {
            int u, v, w;
            cin >> u >> v >> w;
            graph[u].emplace_back(v, w, e);
            graph[v].emplace_back(u, w, e);
      }
      makeTree();
      dfsSize();
      nxt[1] = 1;
      hld();
      depth[1] = 1;
      dfs();
      N = tnode;
 
      for (int i = 1; i <= tnode; i++)
      {
//            cerr << i << ' ' << nodeCost[i] << ' ' << Id[i] << endl;
            updatehld(i, nodeCost[i]);
      }
 
      int q;
      cin >> q;
      while (q--)
      {
            int type;
            cin >> type;
            if (type == 1)
            {
                  int id, val;
                  cin >> id >> val;
                  updatehld(Id[id], val);
            }
            else
            {
                  int u, v;
                  cin >> u >> v;
                  cout << queryhld(u, v) << endl;
            }
      }
}
 

No comments:

Post a Comment

Note: Only a member of this blog may post a comment.