diff --git a/du5/program.c b/du5/program.c
index 3e7f18b..9f55d3b 100644
--- a/du5/program.c
+++ b/du5/program.c
@@ -20,7 +20,7 @@ Tree* readTree() {
 
     if (line[0] == '\0') return NULL;
 
-    Tree *node = (Tree*)malloc(sizeof(Tree));
+    Tree *node = malloc(sizeof(Tree));
     node->yes = NULL;
     node->no = NULL;
 
@@ -41,6 +41,14 @@ Tree* readTree() {
     return node;
 }
 
+int EmptyLine() {
+    char line[SIZE];
+
+    if (fgets(line, SIZE, stdin) == NULL) return 0;
+
+    return (line[0] == '\n' || line[0] == '\r');
+}
+
 int Leaves(Tree *node) {
     if (!node) return 0;
     if (node->isAnswer) return 1;
@@ -57,24 +65,19 @@ void freeTree(Tree *node) {
 void start(Tree *node) {
     char input[SIZE];
 
-    while (node && node->isAnswer == 0) {
+    while (node && !node->isAnswer) {
         printf("%s\n", node->text);
 
         if (fgets(input, SIZE, stdin) == NULL) {
-            return;
-}
-
-        int i = 0;
-        while (input[i] == ' ' || input[i] == '\t') i++;
-
-        if (input[i] == 'a') {
-            node = node->yes;
-        } else if (input[i] == 'n') {
-            node = node->no;
-        } else {
-            printf("Nespravny vstup\n");
+            printf("Koniec vstupu\n");
             return;
         }
+
+        int i = 0;
+        while (input[i] == ' ' || input[i] == '\t' || input[i] == '\r') i++;
+
+        if (input[i] == 'a') node = node->yes;
+        else node = node->no;
     }
 
     if (node && node->isAnswer) {
@@ -87,14 +90,19 @@ int main() {
     printf("Expert z bufetu to vie.\n");
 
     Tree *root = readTree();
+
     if (!root) {
         printf("Chyba nacitania\n");
         return 0;
     }
 
-    int count = Leaves(root);
+    if (!EmptyLine()) {
+        printf("Chyba nacitania\n");
+        freeTree(root);
+        return 0;
+    }
 
-    printf("Pozna %d druhov ovocia a zeleniny.\n", count);
+    printf("Pozna %d druhov ovocia a zeleniny.\n", Leaves(root));
     printf("Odpovedajte 'a' pre prvu moznost alebo 'n' pre druhu moznost.\n");
 
     start(root);