diff --git a/lab2/src/main/java/ru/spbstu/telematics/java/App.java b/lab2/src/main/java/ru/spbstu/telematics/java/App.java index 261d6c7..3af9d9d 100644 --- a/lab2/src/main/java/ru/spbstu/telematics/java/App.java +++ b/lab2/src/main/java/ru/spbstu/telematics/java/App.java @@ -8,5 +8,10 @@ public class App System.out.println(tree.add(10)); System.out.println(tree.add(20)); System.out.println(tree.add(10)); + + System.out.println("Содержимое дерева:"); + for (Integer val: tree) { + System.out.println(val); + } } } diff --git a/lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java b/lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java index 840d7e3..2aaf7c1 100644 --- a/lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java +++ b/lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java @@ -1,7 +1,10 @@ package ru.spbstu.telematics.java; +import java.util.Iterator; +import java.util.NoSuchElementException; +import java.util.Stack; -public class MyTreeSet> { +public class MyTreeSet> implements Iterable { // Класс для представления узла дерева private class Node { E value; @@ -119,4 +122,42 @@ public class MyTreeSet> { } return current; } + + @Override + public Iterator iterator() { + return new MyTreeSetIterator(); + } + + private class MyTreeSetIterator implements Iterator { + // На вершине стека всегда будет лежать наименьший элемент в дереве + Stack nodeStack = new Stack<>(); + + MyTreeSetIterator() { + pushAllLeftToStack(root); + } + + void pushAllLeftToStack(Node current) { + if (current == null) return; + nodeStack.push(current); + pushAllLeftToStack(current.left); + } + + @Override + public boolean hasNext() { + return nodeStack.size() != 0; + } + + @Override + public E next() { + if (!hasNext()) throw new NoSuchElementException(); + + Node minNode = nodeStack.pop(); + + if (minNode.right != null) { + pushAllLeftToStack(minNode.right); + } + + return minNode.value; + } + } } diff --git a/lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java b/lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java index 1f6468f..15daa40 100644 --- a/lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java +++ b/lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java @@ -4,6 +4,8 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import static org.junit.jupiter.api.Assertions.*; +import java.util.Iterator; +import java.util.NoSuchElementException; import java.util.TreeSet; class MyTreeSetTests { @@ -71,4 +73,27 @@ class MyTreeSetTests { assertEquals(myTreeSet.remove(50), treeSet.remove(50)); assertEquals(myTreeSet.size(), treeSet.size()); } + + @Test + void testIterator() { + treeSet.add(15); + myTreeSet.add(15); + treeSet.add(250); + myTreeSet.add(250); + treeSet.add(50); + myTreeSet.add(50); + Iterator itMyTreeSet = myTreeSet.iterator(); + Iterator itTreeSet = treeSet.iterator(); + + assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext()); + assertEquals(itMyTreeSet.next(), itTreeSet.next()); + assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext()); + assertEquals(itMyTreeSet.next(), itTreeSet.next()); + assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext()); + assertEquals(itMyTreeSet.next(), itTreeSet.next()); + + assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext()); + assertThrows(NoSuchElementException.class, () -> itMyTreeSet.next()); + assertThrows(NoSuchElementException.class, () -> itTreeSet.next()); + } }