一道關於二叉樹的位元組面試題的思考

琴水玉發表於2020-10-27

技術人的精神,就是追根究底,把一個事情徹底弄清楚吧!


題目

眾所周知,位元組在一二面的末尾,會隨機抽一道演算法題,當場寫程式碼。我抽到的題目如下:

二叉樹根節點到葉子節點的所有路徑和。給定一個僅包含數字 0−9 的二叉樹,每一條從根節點到葉子節點的路徑都可以用一個數字表示。例如根節點到葉子節點的一條路徑是 1→2→3,那麼這條路徑就用 123 來代替。找出根節點到葉子節點的所有路徑表示的數字之和。

例如:這棵二叉樹一共有兩條路徑,根節點到左葉子節點的路徑 12 代替,根節點到右葉子節點的路徑用 13 代替。所以答案為12+13=25 。


遞迴解法

看到這個題目,首先想到的,其實就是找到這個二叉樹的從根節點到葉子節點的所有路徑。而要找到所有路徑,第一想到的肯定是遞迴。通過左子樹的遞迴拿到的路徑、右子樹的遞迴拿到的路徑,以及根節點,得出最終的所有路徑。

演算法如下:

STEP1:如果已經是葉子節點,那麼構造一條路徑列表,該路徑只有一個元素即葉子節點的值,然後返回【退出條件】。

STEP2: 遞迴找到左子樹到葉子節點的所有路徑列表。對於每條路徑,將根節點加入,從而得到新的結果路徑,並加入;

STEP3:遞迴找到右子樹到葉子節點的所有路徑列表。對於每條路徑,將根節點加入,從而得到新的結果路徑,並加入;

STEP4: 將左右子樹的所有路徑合併成最終的路徑列表【組合子問題的解】。

有兩點說明下:

  • 由於節點的值只有 0-9,因此,可以直接用字串來表示路徑。如果用 List[Integer] ,更靈活,不過會變成列表的列表,處理起來會有點繞。
  • 構建路徑時,使用的是 StringBuilder 的 append 方法,而不是 insert 方法,因此構造的路徑是逆序的。主要考慮到 insert 方法會導致陣列頻繁移動,效率低。具體可以看 StringBuilder 實現。

遞迴程式碼如下:

public List<Path> findAllPaths(TreeNode root) {
        List<Path> le = new ArrayList<>();
        List<Path> ri = new ArrayList<>();
        if (root != null) {

            if (root.left == null && root.right == null) {
                List<Path> single = new ArrayList<>();
                single.add(new Path(root.val));
                return single;
            }

            if (root.left != null) {
                le = findAllPaths(root.left);
                for (Path p: le) {
                    p.append(root.val);
                }
            }
            if (root.right != null) {
                ri = findAllPaths(root.right);
                for (Path p: ri) {
                    p.append(root.val);
                }
            }
        }
        List<Path> paths = new ArrayList<>();
        paths.addAll(le);
        paths.addAll(ri);
        return paths;
    }


class Path {
    StringBuilder s = new StringBuilder();
    public Path() { }

    public Path(Integer i) {
        s.append(i);
    }

    public Path(List list) {
        list.forEach( e-> {
            s.append(e);
        });
    }

    public Path(String str) { this.s = new StringBuilder(str); }

    public Long getValue() {
        return Long.parseLong(s.reverse().toString());
    }

    public StringBuilder append(Integer i) {
        return s.append(i);
    }

    public String toString() {
        return s.reverse().toString();
    }
}


class TreeNode {

    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }

    public int height() {
        if (left == null && right == null) {
            return 1;
        }
        int leftHeight = 0;
        int rightHeight = 0;
        if (left != null) {
            leftHeight = left.height();
        }
        if (right != null) {
            rightHeight = right.height();
        }
        return 1 + max(leftHeight, rightHeight);
    }
}

關鍵點

實際上,我在面試當場沒有做出來,但在面試後的十分鐘,我就把程式碼寫出來了。可能在面試的時候有點緊張,有個地方一直卡住了。

類似二叉樹、動態規劃的問題,由於有多條分支,從思維上來說,不像處理陣列、連結串列那樣是一種線性思維,而是需要一種非線性思維,因此,多做類似的題目,對思維的鍛鍊是很有益的,—— 能夠幫助人擺脫固有的線性思維。

一般來說,演算法問題,通常可以分為兩步:1. 劃分子問題; 2. 將子問題的解組合成原問題的解。 劃分子問題,相對容易一點,但如果劃分不合理,就難以想清楚如何去組合解。我一開始就想到了要用子樹的解與根節點來組合,但是一直糾結在對求出單條路徑的思考上,而不是把所有路徑作為子問題的解。這樣,我就難以想到如何去組合得到最終解。但面試結束之後,我腦子裡閃過左子樹的所有路徑列表,頓時就明白如何組合了。因此,有時,把“所有”作為子問題的解,再跟上層節點組合,反而能容易地得到原問題的解。此外,遞迴要特別注意退出條件。

推薦可以多做二叉樹、動態規劃的題目,能夠很好地鍛鍊劃分子問題、組合子問題的解來求解的技能。

非遞迴演算法

實現遞迴解法,只是一個開始。遞迴演算法很簡潔,但執行效率很低,而且容易棧溢位。如果一個足夠大的二叉樹,就能讓遞迴程式碼無法執行下去。因此,需要尋求非遞迴實現。

非遞迴實現,往往需要藉助於棧。我們需要模擬一下如何用棧來訪問二叉樹。如下圖所示:

可以先找找規則,往往規則就是程式碼的路徑。

  • 每次走到一個節點,先將節點值入棧;
  • 走到葉子節點時,說明已經走到路徑的盡頭,可以記錄下這條路徑。

第一版非遞迴實現如下。用一個棧來儲存二叉樹的訪問節點。如果是葉子節點,就記錄路徑,然後將葉子節點出棧,繼續訪問。

public List<Path> findAllPathsNonRecDeadLoop(TreeNode root) {

        List<Path> allPaths = new ArrayList<>();
        Stack<Integer> s = new DyStack<Integer>();

        TreeNode p = root;
        while(p != null) {
            s.push(p.val);
            if (p.left == null && p.right == null) {
                allPaths.add(new Path(s.unmodifiedList()));
                s.pop();
                if (s.isEmpty()) {
                    break;
                }
            }
            if (p.left != null) {
                p = p.left;
            }
            else if (p.right != null) {
                p = p.right;
            }
        }
        return allPaths;
    }

不過,這個程式碼實現會陷入死迴圈。為什麼呢?因為它會無止境重複進入左子樹,而且回溯的時候,也沒法找到父節點。

回溯

為了解決死迴圈的問題,我們需要加一些支援:進入某個節點時,必須記下該節點的父節點,以及該父節點是否訪問過左右子樹。這個資訊用 TraceNode 來表示。由於始終需要回溯,因此,TraceNode 必須放在棧中,在適當的時候彈出,就像儲存現場一樣。當遍歷的時候,需要記錄已經訪問的節點,不重複訪問,也需要避免將中間節點重複壓棧。

重新理一下。對於當前節點,有四種情形需要考慮:

  • 當前節點是葉子節點。記錄路徑、出棧 treeData, 出棧 traceNode ,回溯到父節點;
  • 當前節點不是葉子節點,有左子樹,則需要記錄該節點指標及左子樹已訪問,並進入左子樹;
  • 當前節點不是葉子節點,有右子樹,則需要記錄該節點指標及右子樹已訪問,並進入右子樹;
  • 當前節點不是葉子節點,有左右子樹且均已訪問,出棧 treeData, 出棧 traceNode ,回溯到父節點。

第二版的遞迴實現如下:

public List<Path> findAllPathsNonRec(TreeNode root) {

        List<Path> allPaths = new ArrayList<>();
        Stack<Integer> treeData = new DyStack<>();
        Stack<TraceNode> trace = new DyStack<>();

        TreeNode p = root;
        TraceNode traceNode = TraceNode.getNoAccessedNode(p);
        while(p != null) {
            if (p.left == null && p.right == null) {
                // 葉子節點的情形,需要記錄路徑,並回溯到父節點
                treeData.push(p.val);
                allPaths.add(new ListPath(treeData.unmodifiedList()));
                treeData.pop();
                if (treeData.isEmpty()) {
                    break;
                }
                traceNode = trace.pop();
                p = traceNode.getParent();
                continue;
            }
            else if (traceNode.needAccessLeft()) {
                // 需要訪問左子樹的情形
                treeData.push(p.val);
                trace.push(TraceNode.getLeftAccessedNode(p));
                p = p.left;
            }
            else if (traceNode.needAccessRight()) {
                // 需要訪問右子樹的情形
                if (traceNode.hasNoLeft()) {
                    treeData.push(p.val);
                }
                if (!traceNode.hasAccessedLeft()) {
                    // 訪問左節點時已經入棧過,這裡不重複入棧
                    treeData.push(p.val);
                }
                trace.push(TraceNode.getRightAccessedNode(p));
                p = p.right;
                if (p.left != null) {
                    traceNode = TraceNode.getNoAccessedNode(p);
                }
                else if (p.right != null) {
                    traceNode = TraceNode.getLeftAccessedNode(p);
                }
            }
            else if (traceNode.hasAllAccessed()) {
                // 左右子樹都已經訪問了,需要回溯到父節點
                if (trace.isEmpty()) {
                    break;
                }
                treeData.pop();
                traceNode = trace.pop();
                p = traceNode.getParent();
            }
        }
        return allPaths;
    }

class TraceNode {

    private TreeNode parent;
    private int accessed;  // 0 均未訪問 1 已訪問左 2 已訪問右

    public TraceNode(TreeNode parent, int child) {
        this.parent = parent;
        this.accessed = child;
    }

    public static TraceNode getNoAccessedNode(TreeNode parent) {
        return new TraceNode(parent, 0);
    }

    public static TraceNode getLeftAccessedNode(TreeNode parent) {
        return new TraceNode(parent, 1);
    }

    public static TraceNode getRightAccessedNode(TreeNode parent) {
        return new TraceNode(parent, 2);
    }

    public boolean needAccessLeft() {
        return parent.left != null && accessed == 0;
    }

    public boolean needAccessRight() {
        return parent.right != null && accessed < 2;
    }

    public boolean hasAccessedLeft() {
        return parent.left == null || (parent.left != null && accessed == 1);
    }

    public boolean hasNoLeft() {
        return parent.left == null;
    }

    public boolean hasAllAccessed() {
        if (parent.left != null && parent.right == null && accessed == 1) {
            return true;
        }
        if (parent.right != null && accessed == 2) {
            return true;
        }
        return false;
    }

    public TreeNode getParent() {
        return parent;
    }

    public int getAccessed() {
        return accessed;
    }
}

關於是否已訪問左右子樹的判斷都隱藏在 TraceNode 裡,findAllPathsNonRec 方法不感知這個。後續如果覺得用 int 來表示 accessed 空間效率不高,可以內部重構,對 findAllPathsNonRec 無影響。這就是封裝的益處。

測試

遞迴程式碼和非遞迴程式碼都是容易有 BUG 的,需要仔細測試下。測試用例通常至少要包括:

  • C1: 單個根節點樹;
  • C2: 單個根節點 + 左節點;
  • C3: 單個根節點 + 右節點;
  • C4: 單個根節點 + 左右節點;
  • C5: 普通的二叉樹,左右隨機;
  • 複雜的二叉樹,非常大。

如何構造複雜的二叉樹呢?可以採用構造法。基於簡單的 C2,C3,C4,將一棵樹的根節點連線到另一棵樹的左葉子節點或右葉子節點上。複雜結構總是由簡單結構來組合而成。

測試程式碼如下。用 TreeBuilder 註解來表示構造的二叉樹,從而能夠批量拿到這些方法構造的樹,進行測試。

   public static void main(String[] args) {
        TreePathSum treePathSum = new TreePathSum();
        Method[] methods = treePathSum.getClass().getDeclaredMethods();
        for (Method m: methods) {
            if (m.isAnnotationPresent(TreeBuilder.class)) {
                try {
                    TreeNode t = (TreeNode) m.invoke(treePathSum, null);
                    System.out.println("height: " + t.height());
                    treePathSum.test2(t);
                } catch (Exception ex) {
                    System.err.println(ex.getMessage());
                }

            }
        }
    }

    public void test(TreeNode root) {

        System.out.println("Rec Implementation");

        List<Path> paths = findAllPaths(root);
        Long sum = paths.stream().collect(Collectors.summarizingLong(Path::getValue)).getSum();
        System.out.println(paths);
        System.out.println(sum);

        System.out.println("Non Rec Implementation");

        List<Path> paths2 = findAllPathsNonRec(root);
        Long sum2 = paths2.stream().collect(Collectors.summarizingLong(Path::getValue)).getSum();
        System.out.println(paths2);
        System.out.println(sum2);

        assert sum == sum2;
    }

    public void test2(TreeNode root) {
        System.out.println("Rec Implementation");
        List<Path> paths = findAllPaths(root);
        System.out.println(paths);

        System.out.println("Non Rec Implementation");
        List<Path> paths2 = findAllPathsNonRec(root);
        System.out.println(paths2);

        assert paths.size() == paths2.size();
        for (int i=0; i < paths.size(); i++) {
            assert paths.get(i).toString().equals(paths2.get(i).toString());
        }

    }

    @TreeBuilder
    public TreeNode buildTreeOnlyRoot() {
        TreeNode tree = new TreeNode(9);
        return tree;
    }

    @TreeBuilder
    public TreeNode buildTreeWithL() {
        return buildTreeWithL(5, 1);
    }

    public TreeNode buildTreeWithL(int rootVal, int leftVal) {
        TreeNode tree = new TreeNode(rootVal);
        TreeNode left = new TreeNode(leftVal);
        tree.left = left;
        return tree;
    }

    @TreeBuilder
    public TreeNode buildTreeWithR() {
        return buildTreeWithR(5,2);
    }

    public TreeNode buildTreeWithR(int rootVal, int rightVal) {
        TreeNode tree = new TreeNode(rootVal);
        TreeNode right = new TreeNode(rightVal);
        tree.right = right;
        return tree;
    }

    @TreeBuilder
    public TreeNode buildTreeWithLR() {
        return buildTreeWithLR(5,1,2);
    }

    public TreeNode buildTreeWithLR(int rootVal, int leftVal, int rightVal) {
        TreeNode tree = new TreeNode(rootVal);
        TreeNode left = new TreeNode(leftVal);
        TreeNode right = new TreeNode(rightVal);
        tree.right = right;
        tree.left = left;
        return tree;
    }

    Random rand = new Random(System.currentTimeMillis());

    @TreeBuilder
    public TreeNode buildTreeWithMore() {
        TreeNode tree = new TreeNode(5);
        TreeNode left = new TreeNode(1);
        TreeNode right = new TreeNode(2);
        TreeNode left2 = new TreeNode(3);
        TreeNode right2 = new TreeNode(4);
        tree.right = right;
        tree.left = left;
        left.left = left2;
        left.right = right2;
        return tree;
    }

    @TreeBuilder
    public TreeNode buildTreeWithMore2() {
        TreeNode tree = new TreeNode(5);
        TreeNode left = new TreeNode(1);
        TreeNode right = new TreeNode(2);
        TreeNode left2 = new TreeNode(3);
        TreeNode right2 = new TreeNode(4);
        tree.right = right;
        tree.left = left;
        right.left = left2;
        right.right = right2;
        return tree;
    }

    public TreeNode treeWithRandom() {
        int c = rand.nextInt(3);
        switch (c) {
            case 0: return buildTreeWithL(rand.nextInt(9), rand.nextInt(9));
            case 1: return buildTreeWithR(rand.nextInt(9), rand.nextInt(9));
            case 2: return buildTreeWithLR(rand.nextInt(9), rand.nextInt(9), rand.nextInt(9));
            default: return buildTreeOnlyRoot();
        }
    }

    public TreeNode linkRandom(TreeNode t1, TreeNode t2) {
        if (t2.left == null) {
            t2.left = t1;
        }
        else if (t2.right == null) {
            t2.right = t1;
        }
        else {
            int c = rand.nextInt(4);
            switch (c) {
                case 0: t2.left.left = t1;
                case 1: t2.left.right = t1;
                case 2: t2.right.left = t1;
                case 3: t2.right.right = t1;
                default: t2.left.left = t1;
            }
        }
        return t2;
    }

    @TreeBuilder
    public TreeNode buildTreeWithRandom() {
        TreeNode root = treeWithRandom();
        int i = 12;
        while (i > 0) {
            TreeNode t = treeWithRandom();
            root = linkRandom(root, t);
            i--;
        }
        return root;
    }

經測試,發現第二版非遞迴程式在某種情況下,還是有 BUG 。這說明某些基本情形還是沒覆蓋到。用如下測試用例除錯,發現就有問題:

@TreeBuilder
    public TreeNode buildTreeWithMore4() {
        TreeNode tree = new TreeNode(5);
        TreeNode left = new TreeNode(1);
        TreeNode right = new TreeNode(2);
        TreeNode left2 = new TreeNode(3);
        TreeNode right2 = new TreeNode(4);
        TreeNode right3 = new TreeNode(6);
        tree.right = right;
        tree.left = left;
        left.right = right3;
        right.right = left2;
        left2.right = right2;
        return tree;
    }

回溯再思考

問題在哪裡?當初次進入沒有左子樹的右子樹時,會有問題。這說明,我還沒有真正弄明白整個回溯過程。重新再理一下回溯過程:

  • 有一個用來指向當前訪問節點的指標 p ;
  • 有一個用來儲存已訪問節點值的棧 treeData;
  • 有一個用來回溯的儲存最近一次訪問的節點資訊的棧 trace ;
  • 有一個用來指明往哪個方向走的 traceNode 。

問題在於我沒想清楚 traceNode 到底是什麼含義。 traceNode 的 parent 和 accessed 到底該存放什麼。實際上,traceNode 和 p 是配套使用的。p 是當前進入的節點的指標,而 traceNode 用來指明進入 p 之後,該往哪裡走。 traceNode 的來源應該有兩個:

  • 第一次進入 p 時,這時候,左右子樹都沒有訪問過,parent 應該與 p 相同,而 accessed 總是初始化為 0 ;
  • 訪問了 p 的左子樹或右子樹,回溯進入 p 時,這時候 parent 應該是 p 的父節點,從 trace 裡拿到。

第二版非遞迴程式正是沒有考慮到第一次進入 p 時的情況。 如下程式碼所示。當 p 進入左子樹時,需要將最近一次的父節點資訊入棧 trace ,同時需要將 traceNode 設定為初始進入 p 時的情形。進入右子樹類似。這一點正是第二版非遞迴程式沒有想清楚的地方。

trace.push(TraceNode.getLeftAccessedNode(p));
p = p.left;
traceNode = TraceNode.getNoAccessedNode(p);

我們做一些修改,得到了第三版非遞迴程式。經測試是 OK 的。

public List<Path> findAllPathsNonRec(TreeNode root) {

        List<Path> allPaths = new ArrayList<>();
        Stack<Integer> treeData = new DyStack<>();
        Stack<TraceNode> trace = new DyStack<>();

        TreeNode p = root;
        TraceNode traceNode = TraceNode.getNoAccessedNode(p);
        while(p != null) {
            if (p.left == null && p.right == null) {
                // 葉子節點的情形,需要記錄路徑,並回溯到父節點
                treeData.push(p.val);
                allPaths.add(new ListPath(treeData.unmodifiedList()));
                treeData.pop();
                if (treeData.isEmpty()) {
                    break;
                }
                traceNode = trace.pop();
                p = traceNode.getParent();
                continue;
            }
            else if (traceNode.needAccessLeft()) {
                // 需要訪問左子樹的情形
                treeData.push(p.val);
                trace.push(TraceNode.getLeftAccessedNode(p));
                p = p.left;
                traceNode = TraceNode.getNoAccessedNode(p);
            }
            else if (traceNode.needAccessRight()) {
                // 需要訪問右子樹的情形
                if (traceNode.hasNoLeft()) {
                    treeData.push(p.val);
                }
                if (!traceNode.hasAccessedLeft()) {
                    // 訪問左節點時已經入棧過,這裡不重複入棧
                    treeData.push(p.val);
                }
                trace.push(TraceNode.getRightAccessedNode(p));
                p = p.right;
                traceNode = TraceNode.getNoAccessedNode(p);
            }
            else if (traceNode.hasAllAccessed()) {
                // 左右子樹都已經訪問了,需要回溯到父節點
                if (trace.isEmpty()) {
                    break;
                }
                treeData.pop();
                traceNode = trace.pop();
                p = traceNode.getParent();
            }
        }
        return allPaths;
    }

優化

擴充套件性

由於題目中所給的節點值為 0-9, 因此,前面取巧用了字串來表示路徑。如果節點值不為 0-9 呢?如果依然要用字串表示,則需要分隔符。現在,我們用列表來表示路徑。封裝的好處,就在於可以替換實現,而儘量少地改變客戶端程式碼(在這裡是findAllPaths 和 findAllPathsNonRec 方法)。

這裡,Path 類改成介面,原來的 Path 類改成 StringPath ,然後用 StringPath 替換 Path 。 將原來 StringPath 用到的方法,定義成介面方法。只用到了 append 和 getValue 方法。不過,構造器方法引數也要相容。這樣,只要把原來的 StringPath 改成 ListPath ,其它基本不用動,就可以執行通過。

interface Path {
    void append(Integer i);
    Long getValue();
}

class StringPath implements Path { // code as before }

class ListPath implements Path {
    List<Integer> path = new ArrayList<>();

    public ListPath(int i) {
        this.path.add(i);
    }

    public ListPath(List list) {
        this.path.addAll(list);
    }

    @Override
    public void append(Integer i) {
        path.add(i);
    }

    @Override
    public Long getValue() {
        StringBuilder s = new StringBuilder();
        path.forEach( e-> {
            s.append(e);
        });
        return Long.parseLong(s.reverse().toString());
    }

    public String toString() {
        return StringUtils.join(path.toArray(), "");
    }
}

小結

花了一天弄懂二叉樹回溯的玩法。技術人的精神,就是追根究底,把一個事情徹底弄清楚吧!

在本文中,我們通過一個二叉樹的路徑尋找面試題,討論了遞迴和非遞迴解法,探討了非遞迴過程中遇到的問題,模擬了二叉樹的回溯,對於理解二叉樹的訪問是很有益的。而對於回溯演算法的理解,鍛鍊了非線性思維。此外,當程式有 BUG 時,往往是某個方面沒想得足夠明白導致。堅持思考,清晰定義,就能向正確再邁進一步。

不看答案,自己弄明白一個問題,收穫大大的!

本文完整原始碼見:ALLIN 工程裡:zzz.study.datastructure.tree.TreePathSum


相關文章