Heavy Light decomposition (Java)

백준 13510 트리와 쿼리 1

참고 코드

import static java.lang.Integer.*;

import java.io.*;
import java.util.*;


class Segment {

    int[] tree;

    public Segment(int size) {
        this.tree = new int[size];
    }

    void update(int node , int start , int end , int index , int value) {
        if (index>end || index<start) {
            return;
        }

        if (start==end) {
            tree[node] = value;
            return;
        }
        int mid = (start+end)/2;
        update(node<<1 , start , mid , index , value);
        update(node<<1|1 , mid+1 , end , index , value);
        tree[node] = Math.max(tree[node<<1] , tree[node<<1|1]);

    }

    int query(int node , int start , int end , int left , int right) {
        if (left>end || right<start) {
            return 0;
        }

        if (left<=start && right>=end) {
            return tree[node];
        }
        int mid = (start+end)/2;

        return Math.max( query(node<<1 , start , mid , left , right) , query(node<<1|1 , mid+1 , end , left ,right));
    }
}

class HLD {

    /**
     * Heavy Light Decomposition
     */

    Segment segment;
    int LENGTH = 100001;
    int[] subtreeSize = new int[LENGTH];
    int[] depth  = new int[LENGTH];
    int[] parent  = new int[LENGTH];
    int[] top  = new int[LENGTH];
    int[] in  = new int[LENGTH];
    int[] out  = new int[LENGTH];
    boolean[] visited = new boolean[LENGTH];
    List<Integer>[] graph= new List[LENGTH], child = new List[LENGTH];
    List<int[]> edge = new ArrayList<>();


    int cnt = 0; // use for ett
    int ans; // use for query
    int MAX;

    public HLD(int size) {
        segment = new Segment(4*size);
        MAX = size;
        for (int i = 0; i < LENGTH; i++) {
            graph[i] = new ArrayList<>();
            child[i] = new ArrayList<>();
        }
    }


    void dfs1(int node) {
        visited[node] = true;
        for (Integer nextNode : graph[node]) {
            if (visited[nextNode]) {
                continue;
            }
            visited[nextNode] = true;
            child[node].add(nextNode);
            dfs1(nextNode);
        }
    }

    void dfs2(int node) {
        subtreeSize[node] = 1;
        for (int i = 0 ; i < child[node].size() ; i++) {
            int nextNode = child[node].get(i);
            depth[nextNode] = depth[node] + 1;
            parent[nextNode] = node;

            dfs2(nextNode);
            subtreeSize[node] += subtreeSize[nextNode];
            if (subtreeSize[nextNode] > subtreeSize[child[node].get(0)]) {
                int temp = child[node].get(0);
                child[node].set(0, nextNode); // swap , child[node].set(0, nextNode) 만 하면 안됨.
                child[node].set(i , temp);  // swap(child[node].get(0) , child[node].get(i) )
            }
        }
    }

    void ett(int node) {
        in[node] = ++cnt;
        for (Integer nextNode : child[node]) {
            top[nextNode] = nextNode.equals(child[node].get(0)) ? top[node] : nextNode;
            ett(nextNode);
        }
        out[node] = cnt;
    }

    void update(int u , int v , int w) {
        if (depth[u] < depth[v]) {
            int temp = u;
            u = v;
            v = temp;
        }
        segment.update(1 , 1 , MAX , in[u] , w);
    }

    int query(int x, int y) {
        ans = 0;
        while (top[x] != top[y]) {
            if (depth[top[x]] < depth[top[y]]) {
                int temp = x;
                x = y;
                y = temp;
            }
            int start = top[x];
            ans = Math.max(ans, segment.query(1, 1, MAX, in[start], in[x]));
            x = parent[start];
        }

        if (depth[x] > depth[y]) {
            int temp = x;
            x = y;
            y = temp;
        }
        ans = Math.max(ans, segment.query(1, 1, MAX, in[x] + 1, in[y]));
        return ans;
    }
}

public class main91 {


    static int N, u, v, w, M, value , query , a , b;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;

        N = parseInt(br.readLine());

        HLD hld = new HLD(N);
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());
            w = parseInt(st.nextToken());

            hld.graph[u].add(v);
            hld.graph[v].add(u);
            hld.edge.add(new int[]{u,v, w});
        }

        hld.dfs1(1); hld.dfs2(1); hld.ett(1);
        for (int[] ints : hld.edge) {
            hld.update(ints[0], ints[1], ints[2]);
        }

        M = parseInt(br.readLine());
        while (M-->0) {
            st = new StringTokenizer(br.readLine());
            query = parseInt(st.nextToken());
            a = parseInt(st.nextToken());
            b = parseInt(st.nextToken());

            if (query == 1) {
                hld.update(hld.edge.get(a - 1)[0] , hld.edge.get(a-1)[1] , b);
            } else {
                if (a==b) {
                    bw.write("0\n");
                } else {
                    bw.write(Integer.toString(hld.query(a,b))+'\n');
                }
            }
        }

        bw.flush();
    }
}