diff --git a/Graph Theory/Diameter_of_tree.cpp b/Graph Theory/Diameter_of_tree.cpp
index 26656db08..b0d89412b 100644
--- a/Graph Theory/Diameter_of_tree.cpp	
+++ b/Graph Theory/Diameter_of_tree.cpp	
@@ -1,85 +1,62 @@
-#include <bits/stdc++.h>
-using namespace std;
- 
-#define INF 0x3f3f3f3f
-#define MOD 1000000007
 
-#define ll long long
-#define pb push_back
-#define nl printf("\n");
-#define vint vector<int>
-
-vector<pair<int,int>>g[100005];
-int n;
-int dist[100005];
-void shortest_path(int s)
-{
-	for(int i=0;i<100001;i++)
-		dist[i]=INF;
-	set<pair<int,int>> setds;
-
-	setds.insert({0,s});
-	dist[s]=0;
-	while(!setds.empty())
-	{
-		pair<int,int> temp;
-		temp=*(setds.begin());
-		setds.erase(setds.begin());
-
-		int u=temp.second;
-		vector<pair<int,int>>::iterator it;
-		for(it=g[u].begin();it!=g[u].end();it++)
-		{
-			int v = (*it).first;
-			int weight = (*it).second;
-			if(dist[v] > dist[u] + weight)
-			{
-				if(dist[v] != INF)
-					setds.erase(setds.find({dist[v],v}));
-
-				dist[v]=dist[u] + weight;
-				setds.insert({dist[v],v});
-			}
-		}
-	}
-
-	// printf("Vertex   Distance from Source\n");
- //    for (int i = 0; i <= n ; ++i)
- //        printf("%d \t\t %d\n", i, dist[i]);
-
-}
-
-int main()
-{
-    ifstream myFile("task.in");
-    if(!myFile.fail())
-    {
-        assert(freopen("task.in", "r", stdin));
-    }
-    int m;
-	cin>>n>>m;
-	n--;				//starting from 0
-    int total=0;
-    for(int i=0;i<m;i++)
-    {
-    	int x,y,w;
-    	cin>>x>>y>>w;
-    	g[x].pb({y,w});
-    	g[y].pb({x,w});
-    	total+=w;
-    }
-    shortest_path(0);
-    int ma=0,ind; 
-    for(int i=0;i<n;i++)
-   		if(dist[i] > ma) 
-   			ind=i,ma=dist[i];
-   	for(int i=0;100001;i++)
-   		dist[i]=INF;
-   	shortest_path(ind);
-   	for(int i=0;i<n;i++)
-   		if(dist[i] > ma) 
-   			ind=i,ma=dist[i];
-   	cout<<(total - dist[ind]);
-   	nl
-	return 0;
-}
+// Simple C++ program to find diameter 
+// of a binary tree. 
+#include <bits/stdc++.h> 
+using namespace std; 
+  
+/* Tree node structure used in the program */
+struct Node { 
+    int data; 
+    Node* left, *right; 
+}; 
+  
+/* Function to find height of a tree */
+int height(Node* root, int& ans) 
+{ 
+    if (root == NULL) 
+        return 0; 
+  
+    int left_height = height(root->left, ans); 
+  
+    int right_height = height(root->right, ans); 
+  
+    // update the answer, because diameter of a 
+    // tree is nothing but maximum value of 
+    // (left_height + right_height + 1) for each node 
+    ans = max(ans, 1 + left_height + right_height); 
+  
+    return 1 + max(left_height, right_height); 
+} 
+  
+/* Computes the diameter of binary tree with given root. */
+int diameter(Node* root) 
+{ 
+    if (root == NULL) 
+        return 0; 
+    int ans = INT_MIN; // This will store the final answer 
+    int height_of_tree = height(root, ans); 
+    return ans; 
+} 
+  
+struct Node* newNode(int data) 
+{ 
+    struct Node* node = new Node; 
+    node->data = data; 
+    node->left = node->right = NULL; 
+  
+    return (node); 
+} 
+  
+// Driver code 
+int main() 
+{ 
+    struct Node* root = newNode(1); 
+    root->left = newNode(2); 
+    root->right = newNode(3); 
+    root->left->left = newNode(4); 
+    root->left->right = newNode(5); 
+  
+    printf("Diameter is %d\n", diameter(root)); 
+  
+    return 0; 
+}