ปูพื้นฐาน

DP บนต้นไม้: เขียนให้รอดต้นไม้ลึกแสนชั้น โดยไม่พึ่งสแต็กของระบบ

ท่าที่ทุกคนเขียนคือฟังก์ชันเรียกตัวเอง ซึ่งตกรอบเงียบ ๆ เมื่อต้นไม้ลึกเป็นแสนชั้น ที่นี่เปลี่ยนไปใช้โครงสิบบรรทัดที่ลอกไปใช้ได้ทุกข้อ แล้วต่อด้วยคำถามที่ต้องตอบก่อนเขียนสูตรว่ากิ่งหนึ่งกิ่งต้องส่งค่าขึ้นมากี่ค่า พร้อมโจทย์ฝึก 3 ข้อ

อาการ · ตัดสินใจที่โหนดนี้ไม่ได้ ถ้ายังไม่รู้ว่าข้างล่างเป็นยังไง

โจทย์บนต้นไม้ส่วนใหญ่มีหน้าตาเหมือนกันหมด คือคำถามที่ถามถึงโหนดหนึ่งโหนดจะตอบไม่ได้เลย จนกว่าจะรู้คำตอบของทุกกิ่งที่ห้อยอยู่ใต้มัน เช่นกิ่งนี้มีกี่โหนด กิ่งนี้ทางลงยาวสุดเท่าไร หรืออย่างข้อ โมบาย ที่ถามว่ากิ่งนี้เป็นทรงไหน

ท่ามาตรฐานที่ทุกคนเขียนคือฟังก์ชันที่เรียกตัวเอง ลงไปหาลูกก่อน ลูกตอบกลับมาแล้วค่อยรวมคำตอบ ถูกต้องและอ่านง่ายที่สุด และตกในสนามแข่งบ่อยกว่าที่คิด ด้วยเหตุผลที่ไม่เกี่ยวกับความถูกต้องเลย

บทนี้จึงว่าด้วยสองเรื่องที่มาคู่กันเสมอ หนึ่งคือรูปแบบการคิดของ DP บนต้นไม้ และสองคือวิธีเดินต้นไม้โดยไม่เรียกซ้ำ ซึ่งไม่ได้ยากกว่ากันเลย แต่ปลอดภัยกว่ามาก

ลองเอง · ตอบขนาดของกิ่งให้ครบทุกปม

ลองไล่ตอบจากปมล่างขึ้นบนดู แล้วสังเกตว่าคำตอบของปมบนใช้คำตอบของปมล่างที่เพิ่งตอบไปได้พอดี

กุญแจ · ลำดับที่ลูกเสร็จก่อนพ่อ

สิ่งเดียวที่การเรียกซ้ำให้เรา คือการรับประกันว่าตอนที่เราคิดค่าของโหนด v ลูกของมันคิดเสร็จหมดแล้ว ถ้าเราหาลำดับแบบนั้นมาได้ด้วยวิธีอื่น ก็ไม่ต้องเรียกซ้ำอีกต่อไป

วิธีหาง่ายกว่าที่คิด เดินต้นไม้ด้วยสแต็กที่ประกาศเอง แล้วจดว่าเราหยิบโหนดไหนออกมาก่อนหลัง ลำดับที่ได้มีคุณสมบัติว่าพ่ออยู่ก่อนลูกเสมอ เพราะลูกจะถูกใส่เข้าสแต็กก็ต่อเมื่อพ่อถูกหยิบออกมาแล้ว ดังนั้นอ่านลำดับนั้นย้อนกลับ จะได้ลูกก่อนพ่อ ซึ่งคือสิ่งที่ DP ต้องการพอดี

// เดินต้นไม้ด้วยสแต็กที่ประกาศเอง ไม่มีการเรียกซ้ำเลยสักบรรทัด
vector<int> par(n + 1, 0), order;
order.reserve(n);
vector<char> seen(n + 1, 0);

vector<int> st{1};
seen[1] = 1;
while (!st.empty()) {
    int v = st.back();
    st.pop_back();
    order.push_back(v);                 // พ่อถูกใส่ก่อนลูกเสมอ
    for (int u : adj[v])
        if (!seen[u]) { seen[u] = 1; par[u] = v; st.push_back(u); }
}

// อ่าน order ย้อนกลับ จะได้ลูกก่อนพ่อ ซึ่งคือลำดับที่ DP บนต้นไม้ต้องการ
for (int i = n - 1; i >= 0; i--) {
    int v = order[i];
    // ตรงนี้ลูกทุกตัวของ v คิดเสร็จแล้วแน่นอน
}

ตัวอย่างแรก · ขนาดของกิ่ง

ค่าที่ง่ายที่สุดบนต้นไม้คือขนาดของกิ่ง นิยามคือตัวมันเองหนึ่ง บวกขนาดของกิ่งลูกทุกตัว

เขียนเป็นโค้ดแล้วสั้นจนน่าตกใจ เพราะแทนที่จะวนหาลูกของแต่ละโหนด เราให้ลูกเป็นคนบวกขึ้นไปหาพ่อ ระหว่างที่ไล่ลำดับย้อนกลับ ทุกโหนดจึงถูกแตะแค่ครั้งเดียว

ทุกครั้งที่ถึงคิวโหนดที่มีลูก ลูกของมันมีตัวเลขอยู่แล้วทุกตัว นั่นคือคำรับประกันที่การเรียกซ้ำเคยให้เรา และตอนนี้เราได้มันมาจากลำดับแทน โดยไม่ต้องเสียสแต็กเลยสักไบต์

โจทย์ฝึก · ไล่จากง่ายไปยาก

สามข้อนี้ตอบคนละคำถาม ข้อแรกถามว่า "โครงที่ให้ไปเขียนจริงหน้าตายังไง" ข้อสองถามว่า "ถ้าคำตอบไม่ได้อยู่ที่ราก แต่ซ่อนอยู่กลางต้นไม้" ข้อสามถามว่า "ถ้าลูกต้องส่งขึ้นมามากกว่าหนึ่งค่า" ลองคิดเองก่อนแล้วค่อยแตะดูเฉลยครับ

ฝึกข้อ 1 ★☆☆☆☆ · ขนาดของทุกกิ่ง

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    scanf("%d", &n);
    vector<vector<int>> adj(n + 1);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        scanf("%d %d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    vector<int> par(n + 1, 0), order;
    order.reserve(n);
    vector<char> seen(n + 1, 0);
    vector<int> st{1};
    seen[1] = 1;
    while (!st.empty()) {
        int v = st.back(); st.pop_back();
        order.push_back(v);
        for (int u : adj[v])
            if (!seen[u]) { seen[u] = 1; par[u] = v; st.push_back(u); }
    }

    vector<int> sz(n + 1, 1);
    // เริ่มที่ i = n-1 แล้วหยุดที่ 1 เพราะรากไม่มีพ่อให้บวกขึ้นไป
    for (int i = n - 1; i >= 1; i--) sz[par[order[i]]] += sz[order[i]];

    for (int i = 1; i <= n; i++) printf("%d%c", sz[i], i == n ? '\n' : ' ');
    return 0;
}

ฝึกข้อ 2 ★★☆☆☆ · ทางที่ยาวที่สุดในต้นไม้

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    scanf("%d", &n);
    vector<vector<int>> adj(n + 1);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        scanf("%d %d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    vector<int> par(n + 1, 0), order;
    order.reserve(n);
    vector<char> seen(n + 1, 0);
    vector<int> st{1};
    seen[1] = 1;
    while (!st.empty()) {
        int v = st.back(); st.pop_back();
        order.push_back(v);
        for (int u : adj[v])
            if (!seen[u]) { seen[u] = 1; par[u] = v; st.push_back(u); }
    }

    vector<int> down(n + 1, 0);       // ทางที่ยาวที่สุดจาก v ลงไปหาใบในกิ่งของตัวเอง
    int ans = 0;
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        int b1 = 0, b2 = 0;           // สองกิ่งที่ยาวที่สุดของ v
        for (int u : adj[v]) {
            if (u == par[v]) continue;
            int d = down[u] + 1;
            if (d > b1) { b2 = b1; b1 = d; }
            else if (d > b2) b2 = d;
        }
        down[v] = b1;
        ans = max(ans, b1 + b2);      // ทางที่หักมุมที่ v พอดี
    }
    printf("%d\n", ans);
    return 0;
}

ฝึกข้อ 3 ★★★☆☆ · เลือกโหนดที่ห้ามติดกัน

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    scanf("%d", &n);
    vector<long long> w(n + 1);
    for (int i = 1; i <= n; i++) scanf("%lld", &w[i]);
    vector<vector<int>> adj(n + 1);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        scanf("%d %d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    vector<int> par(n + 1, 0), order;
    order.reserve(n);
    vector<char> seen(n + 1, 0);
    vector<int> st{1};
    seen[1] = 1;
    while (!st.empty()) {
        int v = st.back(); st.pop_back();
        order.push_back(v);
        for (int u : adj[v])
            if (!seen[u]) { seen[u] = 1; par[u] = v; st.push_back(u); }
    }

    vector<long long> take(n + 1, 0), skip(n + 1, 0);
    for (int i = n - 1; i >= 0; i--) {
        int v = order[i];
        take[v] += w[v];
        for (int u : adj[v]) {
            if (u == par[v]) continue;
            take[v] += skip[u];                  // เลือก v แล้ว ลูกทุกตัวห้ามถูกเลือก
            skip[v] += max(take[u], skip[u]);    // ไม่เลือก v ลูกจะเลือกหรือไม่ก็ได้
        }
    }
    printf("%lld\n", max(take[1], skip[1]));
    return 0;
}

ก้าวถัดไปของท่านี้ · เมื่อโจทย์ถามคำตอบของทุกปม ไม่ใช่แค่ของราก

ท่าทั้งบทนี้ตอบคำถามที่มองจากรากลงล่าง คือค่าของกิ่งที่ห้อยใต้ปมหนึ่ง ซึ่งพอเลือกรากแล้วก็คิดได้ในรอบเดียว แต่มีโจทย์อีกตระกูลที่ถามหนักกว่านั้น คือถาม คำตอบของทุกปมในฐานะที่ปมนั้นเป็นราก ตัวอย่างที่เข้าใจง่ายสุดคือ "ผลรวมระยะทางจากปมนี้ไปทุกปมที่เหลือ" ซึ่งต้องตอบให้ครบทั้ง n ปม

ทางตรงคือย้ายรากแล้วคิดใหม่ทีละปม ซึ่งเป็น n รอบ รอบละ O(n) รวมเป็น O(n²) ที่ n = 200,000 คือสี่หมื่นล้านครั้ง แตะไม่ได้

ท่าที่แก้เรื่องนี้เรียกว่า การย้ายราก (rerooting) และแนวคิดของมันสั้นมาก คำตอบของปมข้างเคียงต่างจากคำตอบของปมนี้แค่นิดเดียว จึงคำนวณจากกันได้ ไม่ต้องคิดใหม่ทั้งใบ

ลองดูกับโจทย์ผลรวมระยะทาง สมมติเรารู้คำตอบของปม p แล้ว และจะย้ายรากไปที่ลูก v ระยะทางของทุกปมเปลี่ยนไปแค่สองแบบ ปมที่อยู่ในกิ่งของ v ทั้ง sz[v] ปม ใกล้ขึ้นหนึ่งก้าวทุกปม ส่วนปมที่เหลือทั้ง n - sz[v] ปม ไกลขึ้นหนึ่งก้าวทุกปม คำตอบใหม่จึงเป็น

ans[v] = ans[p] − sz[v] + (n − sz[v])

ซึ่งเป็นงานคงที่ต่อหนึ่งปม รวมทั้งต้นไม้จึงเป็น O(n) โครงของท่านี้คือสองรอบ รอบแรกไล่จากใบขึ้นหาราก เก็บของที่มองลงได้อย่างเดียว (ขนาดกิ่ง และผลรวมระยะทางที่มองลง) รอบที่สองไล่จากรากลงหาใบ ย้ายรากทีละก้าวด้วยสูตรข้างบน รอบแรกใช้ลำดับย้อนกลับที่บทนี้สอนไว้แล้ว รอบที่สองใช้ลำดับเดิมแบบไม่ย้อน

// Rerooting: ผลรวมระยะทางจากทุกปมไปยังปมอื่นทั้งหมด คิดครบทุกปมในสองรอบ
// อินพุต: n แล้ว n-1 บรรทัด u v  (ต้นไม้ ปมนับจาก 0)
// เอาต์พุต: ans[0..n-1] คือผลรวมระยะทางจากปมนั้นไปทุกปม
#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    if (scanf("%d", &n) != 1) return 0;
    vector<vector<int>> g(n);
    for (int i = 0; i + 1 < n; i++) {
        int u, v;
        scanf("%d %d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    // ลำดับ BFS จากราก 0 ใช้แทนการเรียกซ้ำ กันสแต็กล้นบนต้นไม้ลึก
    vector<int> order, par(n, -1);
    order.reserve(n);
    order.push_back(0);
    for (int h = 0; h < (int)order.size(); h++) {
        int u = order[h];
        for (int v : g[u])
            if (v != par[u]) {
                par[v] = u;
                order.push_back(v);
            }
    }

    vector<long long> sz(n, 1), down(n, 0), ans(n, 0);
    // รอบที่หนึ่ง จากใบขึ้นหาราก เก็บขนาดกิ่งและผลรวมระยะทางที่มองลงเท่านั้น
    for (int i = n - 1; i >= 1; i--) {
        int v = order[i], p = par[v];
        sz[p] += sz[v];
        down[p] += down[v] + sz[v];
    }
    // รอบที่สอง จากรากลงหาใบ ย้ายรากทีละก้าว
    // ก้าวจาก p ไป v มีปม sz[v] ปมที่ใกล้ขึ้นหนึ่ง และ n - sz[v] ปมที่ไกลขึ้นหนึ่ง
    ans[0] = down[0];
    for (int i = 1; i < n; i++) {
        int v = order[i], p = par[v];
        ans[v] = ans[p] - sz[v] + (long long)(n - sz[v]);
    }
    for (int i = 0; i < n; i++) printf("%lld%c", ans[i], i + 1 == n ? '\n' : ' ');
    return 0;
}

// ตัวตรวจอิสระของ rerooting: BFS ใหม่จากทุกปม แล้วบวกระยะทางตรง ๆ
// ไม่มีการย้ายราก ไม่มีสูตรส่วนต่าง จึงไม่ได้ทดสอบข้ออ้างของท่านั้นเลย
#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    if (scanf("%d", &n) != 1) return 0;
    vector<vector<int>> g(n);
    for (int i = 0; i + 1 < n; i++) {
        int u, v;
        scanf("%d %d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for (int s = 0; s < n; s++) {
        vector<int> dist(n, -1);
        dist[s] = 0;
        vector<int> q{s};
        for (int h = 0; h < (int)q.size(); h++)
            for (int v : g[q[h]])
                if (dist[v] == -1) {
                    dist[v] = dist[q[h]] + 1;
                    q.push_back(v);
                }
        long long tot = 0;
        for (int i = 0; i < n; i++) tot += dist[i];
        printf("%lld%c", tot, s + 1 == n ? '\n' : ' ');
    }
    return 0;
}

รูปทั่วไปของท่านี้ยกไปใช้ได้กับค่าอื่นด้วย ไม่ใช่แค่ผลรวมระยะทาง คำถามที่ต้องตอบให้ได้ก่อนใช้คือ ถ้าย้ายรากไปหนึ่งก้าว ค่าที่เราสนใจเปลี่ยนไปเท่าไร และคำนวณส่วนต่างนั้นจากของที่มีอยู่แล้วได้ไหม ถ้าตอบได้ ท่านี้ใช้ได้ทันที ถ้าตอบไม่ได้ ให้กลับไปดูว่าต้องเก็บอะไรเพิ่มในรอบแรก

สรุปบรรทัดเดียว

สิ่งเดียวที่ DP บนต้นไม้ต้องการคือลำดับที่ลูกเสร็จก่อนพ่อ ซึ่งได้มาจากการเดินด้วยสแต็กแล้วอ่านย้อนกลับ ไม่ต้องพึ่งสแต็กของระบบเลย ส่วนคำถามที่ต้องตอบก่อนเขียนสูตรทุกครั้งคือลูกต้องส่งขึ้นมากี่ค่า ถ้าการตัดสินใจของพ่อไปจำกัดลูก ค่าเดียวไม่เคยพอ

ข้อที่ใช้ของในบทนี้ตรง ๆ คือ โมบาย ซึ่งค่าที่กิ่งส่งขึ้นมาไม่ใช่ตัวเลข แต่เป็นป้ายบอกทรงของกิ่งสามแบบ และเป็นข้อที่โซ่ยาวหนึ่งแสนชั้นเป็นอินพุตที่ถูกกติกา