LeetCode N-ary Tree

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()
}
}