算法

DFS

之前写了一个BFS建图的方法，这次分享一个DFS树遍历的方法。BFS的方法简单易懂，在网上查到的很多人用dp的方法做的，这个DFS的方法很麻烦，不过它还是很值得学习的。如果代码写的不够清楚，建议大家做一下LeetCode 863，那道题和这个有一些相似。

用树的方法做，难点在于怎么样从一个子树target Node回推到另一个子树的距离。

时间复杂度

$O(n^2)$

Java 代码

import java.io.*;

public class Main {
    public static void main(String[] args) throws IOException {
        //Read console
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));

        int n = Integer.parseInt(reader.readLine());
        int[] left = new int[n + 1];
        int[] right = new int[n + 1];
        int[] population = new int[n + 1];

        for (int i = 1; i <= n; i++) {
            String[] str = reader.readLine().split(" ");
            population[i] = Integer.parseInt(str[0]);
            left[i] = Integer.parseInt(str[1]);
            right[i] = Integer.parseInt(str[2]);
        }
        reader.close();

        //Get root id
        boolean[] isChild = new boolean[n + 1];
        for (int i = 1; i <= n; i++) {
            isChild[left[i]] = true;
            isChild[right[i]] = true;
        }

        int rootId = -1;
        for (int i = 1; i <= n; i++) {
            if(!isChild[i]) {
                rootId = i;
                break;
            }
        }

        //Get min cost
        Main main = new Main();
        int minCost = main.getMinCost(population, left, right, rootId);

        //Print result
        System.out.println(minCost);
    }

    private int[] population;
    private int[] left;
    private int[] right;
    private int cost;

    public int getMinCost(int[] population, int[] left, int[] right, int rootId) {
        this.population = population;
        this.left = left;
        this.right = right;
        this.cost = 0;
        int minCost = Integer.MAX_VALUE;

        for (int i = 1; i < left.length; i++) {
            this.cost = 0;
            distance(i, rootId);
            minCost = Math.min(minCost, this.cost);
        }

        return minCost;
    }

    //Distance from id to targetId
    private int distance(int targetId, int id) {
        if (id == 0) {
            return -1;
        }

        if (targetId == id) {
            collectCost(id, 0);
            return 0;
        }

        int l = distance(targetId, left[id]);
        int r = distance(targetId, right[id]);

        //If l is greater than or equal to 0, 
        //it means that the target is in left subtree.
        //Collect the cost of current node and
        //the cost of right subtree of current node.
        if (l >= 0) {
            int d = l + 1;//d is the distance for current node
            cost += d * population[id];
            collectCost(right[id], d + 1);
            return d;
        }

        if (r >= 0) {
            int d = r + 1;
            cost += d * population[id];
            collectCost(left[id], d + 1);
            return d;
        }

        return -1;
    }

    //Collect cost for the node including its left and right subtrees
    private void collectCost(int id, int level) {
        if (id == 0) {
            return;
        }

        cost += level * population[id];

        collectCost(left[id], level + 1);
        collectCost(right[id], level + 1);
    }
}