LeetCode N-ary Tree

Posted on Edited on Disqus: Word count in article: 2.4k Reading time ≈ 2 mins.

N-ary Tree Preorder Traversal

1
2
3
4
5
6
7
8
9
10
class Solution {
    func preorder(_ root: Node?) -> [Int] {
    	var stack = [root], ans = [Int]()
        while let node = stack.popLast(), let n = node {
            ans.append(n.val)
            stack += n.children.reversed()
        }
        return ans
    }
}

N-ary Tree Postorder Traversal

1
2
3
4
5
6
7
8
9
10
class Solution {
    func postorder(_ root: Node?) -> [Int] {
    	var stack = [root], ans = [Int]()
        while let node = stack.popLast(), let n = node {
            ans.append(n.val)
            stack += n.children
        }
        return ans.reversed()
    }
}

N-ary Tree Level Order Traversal

1
2
3
4
5
6
7
8
9
10
11
12
class Solution {
    func levelOrder(_ root: Node?) -> [[Int]] {
        guard let root = root else { return [] }
        var cur = [root], ans = [[root.val]]
        while !cur.isEmpty {
            let children = cur.flatMap { $0.children }
            if !children.isEmpty { ans.append( children.map { $0.val } ) }
            cur = children
        }
        return ans
    }
}

Maximum Depth of N-ary Tree

1
2
3
4
5
6
7
8
9
10
class Solution {
    func maxDepth(_ root: Node?) -> Int {
        var cur = root == nil ? [] : [root!], depth = 0
        while !cur.isEmpty {
            cur = cur.flatMap { $0.children }
            depth += 1
        }
        return depth
    }
}

Encode N-ary Tree to Binary Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Codec {
    func encode(_ root: Node?) -> TreeNode? {
        guard let root = root else { return nil }
        let tn = TreeNode(root.val); var p: TreeNode? = tn
        if let n = root.children.first {
            p?.left = encode(n); p = p?.left
            for n in root.children[1...] {
                p?.right = encode(n); p = p?.right
            }
        }
        return tn
    }

    func decode(_ root: TreeNode?) -> Node? {
        guard let root = root else { return nil }
        let n = Node(root.val); var p = root.left
        if let l = decode(p) { n.children.append(l) }
        while let q = p?.right, let r = decode(q) {
            n.children.append(r); p = q
        }
        return n
    }
}

Serialize and Deserialize N-ary Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Codec {
    func serialize(_ root: Node?) -> String {
        guard let root = root else { return "$" }
    	var s = "\(root.val),\(root.children.count),"
        for child in root.children {
            s += serialize(child)
        }
        return s
    }

    func deserialize(_ data: String) -> Node? {
        let arr = data.components(separatedBy: ",").compactMap { Int($0) }; var i = 0
        func deserialize() -> Node? {
            if i >= arr.count { return nil }
            let n = Node(arr[i])
            i += 2; for _ in 0..<arr[i-1] {
                if let a = deserialize() { n.children.append(a) }
            }
            return n
        }
        return deserialize()
    }
}

Translated by gpt-3.5-turbo