BOJ 5916 농장 관리

📝 문제 정보

링크: https://www.acmicpc.net/problem/5916

🧐 관찰 및 접근

트리가 주어진다.
- 두 정점 사이 경로에 1을 더한다.
- 두 정점 사이 경로의 가중치의 합을 계산한다.
나이브하게는 $O(QN)$이겠지만, HLD와 Lazy 세그먼트 트리를 이용해서 $O(Qlog^2N)$정도에 계산할 수 있을 것 같다.

💻 풀이

코드 (C++):

void solve(){
    int N, M; cin >> N >> M;
    Tree tree(N);
    rep(i, 0, N-1){
        int u, v; cin >> u >> v;
        tree.add_edge(u, v);
    }
    tree.build();

    LazySegmentTree seg(N+1);
    while(M--){
        char op; cin >> op;
        int u, v; cin >> u >> v;
        if(op == 'P'){
            vector<pii> segs = tree.hld_path(u, v, false, false);
            for(auto [L, R]: segs) seg.update(L, R, 1);
        }
        else{
            ll ans = 0;
            vector<pii> segs = tree.hld_path(u, v, false, false);
            for(auto [L, R]: segs) ans += seg.query(L, R);
            cout << ans << '\n';
        }
    }
}