The trick here is to name the state correctly, then let the implementation follow. For 1519. Number of Nodes in the Sub-Tree With the Same Label, the solution in this repository is mainly a graph traversal solution.
Guide
What?
Before optimizing anything, pin down what information is still useful after each move. Instead of trying to be clever immediately, read the code as a sequence of questions:
- What state are we keeping?
- How do we move from one state to the next?
- When do we know the answer is already determined?
For this file, the main tools are: graph traversal, stack.
The notes already sitting in the source point us in the right direction:
- TLE
- 59 / 59 test cases passed, but took too long.
Guide
When?
Reach for this shape when a direct simulation would work logically but waste time revisiting the same information. The accepted code reduces that pressure by storing exactly the information that remains useful later.
The important function names to track are dfs, countSubTrees, visited, counter.
Guide
Why?
The point of the implementation is not to make the code longer. It is to avoid doing the same thinking twice.
- The stack stores unfinished context, which is usually the cleanest way to handle nested or monotonic structure.
- The two-dimensional vector is the memory of the solution: grid state, DP state, or adjacency shape.
- The final return is not magic; it is the invariant after the loops or recursion have finished doing their accounting.
Guide
How?
Walk through the solution in this order:
- Read the setup variables first.
- Follow the main loop or recursive helper next.
- Watch where invalid states get skipped.
- Check which value survives to the return statement.
The most important competitive-programming habit here is to trust the invariant. Once the invariant is right, the loops become much less scary.
Guide
Complexity
- Time: O(n) to O(n log n), depending on the dominant loop or data structure operation
- Space: O(n) in the usual case for auxiliary containers or recursion
Guide
C++ Solution
Your submission
The accepted solution
solution.cpp
01//TLE
02//59 / 59 test cases passed, but took too long.
03class Solution {
04public:
05 int dfs(vector<vector<int>>& adjList, string& labels, int cur, char c, vector<vector<int>>& counter){
06 // cout << "cur: " << cur << ", c: " << c << endl;
07 if(counter[cur][c-'a'] != -1){
08 return counter[cur][c-'a'];
09 }
10
11 int count = (labels[cur] == c);
12
13 for(int nei : adjList[cur]){
14 count += dfs(adjList, labels, nei, c, counter);
15 }
16
17 return counter[cur][c-'a'] = count;
18 };
19
20 vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
21 vector<vector<int>> adjList(n);
22 vector<vector<int>> counter(n, vector<int>(26, -1));
23
24 //bi-directional
25 for(vector<int>& edge : edges){
26 adjList[edge[0]].push_back(edge[1]);
27 adjList[edge[1]].push_back(edge[0]);
28 }
29
30 vector<bool> visited(n, false);
31 //one directional, top-down
32 vector<vector<int>> tree(n);
33
34 stack<int> stk;
35 stk.push(0);
36 visited[0] = true;
37
38 while(!stk.empty()){
39 int cur = stk.top(); stk.pop();
40
41 for(int nei : adjList[cur]){
42 if(!visited[nei]){
43 visited[nei] = true;
44 tree[cur].push_back(nei);
45 stk.push(nei);
46 }
47 }
48 }
49
50 // for(int i = 0; i < n; ++i){
51 // cout << i << " : ";
52 // for(int child : tree[i]) cout << child << " ";
53 // cout << endl;
54 // }
55
56 for(int i = 0; i < n; ++i){
57 dfs(tree, labels, i, labels[i], counter);
58 }
59
60 vector<int> ans(n);
61
62 for(int i = 0; i < n; ++i){
63 ans[i] = counter[i][labels[i]-'a'];
64 }
65
66 return ans;
67 }
68};
69
70//DFS
71//Runtime: 1564 ms, faster than 40.21% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
72//Memory Usage: 290.5 MB, less than 100.00% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
73class Solution {
74public:
75 vector<int> dfs(vector<vector<int>>& tree, string& labels, int cur, vector<int>& ans){
76 vector<int> counter(26, 0);
77 counter[labels[cur]-'a'] = 1;
78
79 for(int nei : tree[cur]){
80 vector<int> subcounter = dfs(tree, labels, nei, ans);
81 for(int i = 0; i < counter.size(); ++i){
82 counter[i] += subcounter[i];
83 }
84 }
85 //since we are using dfs, now counter's coverage is the subtree rooted at cur
86 ans[cur] += counter[labels[cur]-'a'];
87
88 return counter;
89 };
90
91 vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
92 vector<vector<int>> adjList(n);
93 vector<vector<int>> counter(n, vector<int>(26, -1));
94
95 //bi-directional
96 for(vector<int>& edge : edges){
97 adjList[edge[0]].push_back(edge[1]);
98 adjList[edge[1]].push_back(edge[0]);
99 }
100
101 vector<bool> visited(n, false);
102 //one directional, top-down
103 vector<vector<int>> tree(n);
104
105 stack<int> stk;
106 stk.push(0);
107 visited[0] = true;
108
109 while(!stk.empty()){
110 int cur = stk.top(); stk.pop();
111
112 for(int nei : adjList[cur]){
113 if(!visited[nei]){
114 visited[nei] = true;
115 tree[cur].push_back(nei);
116 stk.push(nei);
117 }
118 }
119 }
120
121 vector<int> ans(n);
122
123 dfs(tree, labels, 0, ans);
124
125 return ans;
126 }
127};
128
129//DFS, visited
130//https://leetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/discuss/743133/JavaPython-3-DFS-w-brief-explanation-and-analysis.
131//Runtime: 1084 ms, faster than 80.42% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
132//Memory Usage: 289.1 MB, less than 100.00% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
133//time: O(N), space: O(N)
134class Solution {
135public:
136 vector<int> dfs(vector<vector<int>>& adjList, string& labels, vector<bool>& visited, int cur, vector<int>& ans){
137 vector<int> counter(26, 0);
138 counter[labels[cur]-'a'] = 1;
139
140 for(int nei : adjList[cur]){
141 if(!visited[nei]){
142 visited[nei] = true;
143 vector<int> subcounter = dfs(adjList, labels, visited, nei, ans);
144 for(int i = 0; i < counter.size(); ++i){
145 counter[i] += subcounter[i];
146 }
147 }
148 }
149 //since we are using dfs, now counter's coverage is the subtree rooted at cur
150 ans[cur] += counter[labels[cur]-'a'];
151
152 return counter;
153 };
154
155 vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
156 vector<vector<int>> adjList(n);
157 vector<vector<int>> counter(n, vector<int>(26, -1));
158
159 //bi-directional
160 for(vector<int>& edge : edges){
161 adjList[edge[0]].push_back(edge[1]);
162 adjList[edge[1]].push_back(edge[0]);
163 }
164
165 vector<int> ans(n);
166 vector<bool> visited(n, false);
167
168 visited[0] = true;
169 dfs(adjList, labels, visited, 0, ans);
170
171 return ans;
172 }
173};
174
175//DFS, remove "visited"
176//https://leetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/discuss/743133/JavaPython-3-DFS-w-brief-explanation-and-analysis.
177//Runtime: 1476 ms, faster than 49.89% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
178//Memory Usage: 278.8 MB, less than 100.00% of C++ online submissions for Number of Nodes in the Sub-Tree With the Same Label.
179//time: O(N), space: O(N)
180class Solution {
181public:
182 vector<int> dfs(vector<vector<int>>& adjList, string& labels, int parent, int cur, vector<int>& ans){
183 vector<int> counter(26, 0);
184 char c = labels[cur]-'a';
185 counter[c] = 1;
186
187 for(int nei : adjList[cur]){
188 if(nei != parent){
189 vector<int> subcounter = dfs(adjList, labels, cur, nei, ans);
190 for(int i = 0; i < counter.size(); ++i){
191 counter[i] += subcounter[i];
192 }
193 }
194 }
195 //since we are using dfs, now counter's coverage is the subtree rooted at cur
196 ans[cur] = counter[c];
197
198 return counter;
199 };
200
201 vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
202 vector<vector<int>> adjList(n);
203 vector<vector<int>> counter(n, vector<int>(26, -1));
204
205 //bi-directional
206 for(vector<int>& edge : edges){
207 adjList[edge[0]].push_back(edge[1]);
208 adjList[edge[1]].push_back(edge[0]);
209 }
210
211 vector<int> ans(n, 0);
212 int parent = -1;
213
214 dfs(adjList, labels, -1, 0, ans);
215
216 return ans;
217 }
218};Cost
Complexity
Time
O(n) to O(n log n), depending on the dominant loop or data structure operation
Dominated by the main traversal, recursion, or data-structure operations in the code.
Space
O(n) in the usual case for auxiliary containers or recursion
Auxiliary state plus the answer structure where the problem requires one.