tbintree.nim 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. discard """
  2. file: "tbintree.nim"
  3. output: "helloworld99110223"
  4. """
  5. type
  6. TBinaryTree[T] = object # TBinaryTree is a generic type with
  7. # with generic param ``T``
  8. le, ri: ref TBinaryTree[T] # left and right subtrees; may be nil
  9. data: T # the data stored in a node
  10. PBinaryTree*[A] = ref TBinaryTree[A] # type that is exported
  11. proc newNode*[T](data: T): PBinaryTree[T] =
  12. # constructor for a node
  13. new(result)
  14. result.data = data
  15. proc add*[Ty](root: var PBinaryTree[Ty], n: PBinaryTree[Ty]) =
  16. # insert a node into the tree
  17. if root == nil:
  18. root = n
  19. else:
  20. var it = root
  21. while it != nil:
  22. # compare the data items; uses the generic ``cmp`` proc that works for
  23. # any type that has a ``==`` and ``<`` operator
  24. var c = cmp(n.data, it.data)
  25. if c < 0:
  26. if it.le == nil:
  27. it.le = n
  28. return
  29. it = it.le
  30. else:
  31. if it.ri == nil:
  32. it.ri = n
  33. return
  34. it = it.ri
  35. proc add*[Ty](root: var PBinaryTree[Ty], data: Ty) =
  36. # convenience proc:
  37. add(root, newNode(data))
  38. proc find*[Ty2](b: PBinaryTree[Ty2], data: Ty2): bool =
  39. # for testing this needs to be recursive, so that the
  40. # instantiated type is checked for proper tyGenericInst envelopes
  41. if b == nil:
  42. result = false
  43. else:
  44. var c = cmp(data, b.data)
  45. if c < 0: result = find(b.le, data)
  46. elif c > 0: result = find(b.ri, data)
  47. else: result = true
  48. iterator preorder*[T](root: PBinaryTree[T]): T =
  49. # Preorder traversal of a binary tree.
  50. # Since recursive iterators are not yet implemented,
  51. # this uses an explicit stack:
  52. var stack: seq[PBinaryTree[T]] = @[root]
  53. while stack.len > 0:
  54. var n = stack.pop()
  55. while n != nil:
  56. yield n.data
  57. add(stack, n.ri) # push right subtree onto the stack
  58. n = n.le # and follow the left pointer
  59. iterator items*[T](root: PBinaryTree[T]): T =
  60. ## Inorder traversal of the binary tree.
  61. var stack: seq[PBinaryTree[T]] = @[]
  62. var n = root
  63. while true:
  64. while n != nil:
  65. add(stack, n)
  66. n = n.le
  67. if stack.len > 0:
  68. n = stack.pop()
  69. yield n.data
  70. n = n.ri
  71. if stack.len == 0 and n == nil: break
  72. proc debug[T](a: PBinaryTree[T]) =
  73. if a != nil:
  74. debug(a.le)
  75. echo a.data
  76. debug(a.ri)
  77. when isMainModule:
  78. var
  79. root: PBinaryTree[string]
  80. x = newNode("hello")
  81. add(root, x)
  82. add(root, "world")
  83. if find(root, "world"):
  84. for str in items(root):
  85. stdout.write(str)
  86. else:
  87. stdout.writeLine("BUG")
  88. var
  89. r2: PBinaryTree[int]
  90. add(r2, newNode(110))
  91. add(r2, 223)
  92. add(r2, 99)
  93. for y in items(r2):
  94. stdout.write(y)
  95. #OUT helloworld99110223