diff --git a/du6/program.c b/du6/program.c
index 69400c6..1545f25 100644
--- a/du6/program.c
+++ b/du6/program.c
@@ -10,7 +10,7 @@ typedef struct TreeNode {
     struct TreeNode *no_branch;
 } TreeNode;
 
-static TreeNode* read_tree_node() {
+static TreeNode* read_tree_node(int *used_lines) {
     char buffer[MAXLINE];
     memset(buffer, 0, sizeof(buffer));
 
@@ -19,6 +19,8 @@ static TreeNode* read_tree_node() {
         return NULL;
     }
 
+    (*used_lines)++;
+
     if (buffer[0] == '\n' || buffer[0] == '\r') {
         return NULL;
     }
@@ -36,8 +38,8 @@ static TreeNode* read_tree_node() {
         return node;
     }
 
-    node->yes_branch = read_tree_node();
-    node->no_branch = read_tree_node();
+    node->yes_branch = read_tree_node(used_lines);
+    node->no_branch = read_tree_node(used_lines);
 
     return node;
 }
@@ -58,10 +60,7 @@ static int count_leafs(TreeNode *root) {
     int left = count_leafs(root->yes_branch);
     int right = count_leafs(root->no_branch);
 
-    int is_leaf = 0;
-    if (root->yes_branch == NULL && root->no_branch == NULL) {
-        is_leaf = 1;
-    }
+    int is_leaf = (root->yes_branch == NULL && root->no_branch == NULL) ? 1 : 0;
 
     return left + right + is_leaf;
 }
@@ -100,23 +99,6 @@ static void skip_empty_line() {
     }
 }
 
-static int has_more_data() {
-    char buffer[MAXLINE];
-    long pos = ftell(stdin);
-
-    if (fgets(buffer, sizeof(buffer), stdin) == NULL) {
-        return 0;
-    }
-
-    int newline = (buffer[0] == '\n' || buffer[0] == '\r');
-
-    if (!newline) {
-        return 1;
-    }
-
-    return 0;
-}
-
 static void run_dialog(TreeNode *root) {
     if (root == NULL) {
         return;
@@ -151,18 +133,34 @@ static void run_dialog(TreeNode *root) {
 }
 
 int main() {
-    TreeNode *root = read_tree_node();
+    int used_lines = 0;
+    TreeNode *root = read_tree_node(&used_lines);
 
     skip_empty_line();
 
     printf("Expert z bufetu to vie.\n");
 
-    if (root == NULL || !validate_tree(root) || has_more_data()) {
+    if (root == NULL || !validate_tree(root)) {
         printf("Chybna databaza\n");
         free_tree(root);
         return 0;
     }
 
+    int total_lines = used_lines;
+
+    long pos = ftell(stdin);
+    char extra[MAXLINE];
+    if (fgets(extra, sizeof(extra), stdin) != NULL) {
+        int newline = (extra[0] == '\n' || extra[0] == '\r');
+        if (!newline) {
+            printf("Chybna databaza\n");
+            free_tree(root);
+            return 0;
+        }
+    } else {
+        fseek(stdin, pos, SEEK_SET);
+    }
+
     int total = count_leafs(root);
 
     printf("Pozna %d druhov ovocia a zeleniny.\n", total);