Got no idea why runtime error happened
Any idea?
https://codeforces.com/problemset/problem/580/C
UPDATED
Thanks to @cs_tree telling me the main reason why python failed but c++ worked.
Therefore I wrote a BFS version in python and passed.
def go_bfs():
bfs = [(cat[0], 0, 0)]
sums = 0
while bfs:
isLeaf = True
level_cat, v, p = bfs.pop(0)
for j in e[v]:
if j == p:
continue
isLeaf = False
if cat[j] == 0:
bfs.append((0, j, v))
elif level_cat + 1 <= m:
bfs.append((level_cat + 1, j, v))
if isLeaf:
sums += 1
return sums
UPDATED
I wrote the same algorithm in c++ and passed... I guess the problem might be input() in python, but still don't know the reason
C++
#include<map>
#include<iostream>
#include<vector>
using namespace std;
int n, m, tmp, v1, v2;
vector<int> cat;
vector<vector<int>> graph;
int go_child(int node, int path_cat, int parent) {
int m_cat = 0, sums = 0;
bool isLeaf = true;
if(cat[node] == 0) {
m_cat = 0;
} else {
m_cat = path_cat + cat[node];
}
if(m_cat > m) {
return 0;
}
for(int i = 0; i < graph[node].size(); i++) {
if(graph[node][i] == parent) {
continue;
}
isLeaf = false;
sums += go_child(graph[node][i], m_cat, node);
}
if(isLeaf) {
return 1;
}
return sums;
}
int main() {
cin >> n >> m;
for(int i = 0; i < n; i++) {
cin >> tmp;
cat.push_back(tmp);
graph.push_back(vector<int>());
}
for(int i = 0; i < n - 1; i++) {
cin >> v1 >> v2;
graph[v1 - 1].push_back(v2 - 1);
graph[v2 - 1].push_back(v1 - 1);
}
cout << go_child(0, 0, -1) << endl;
return 0;
}
Python
def go_child(node, path_cat, parent):
# check if too more cat
m_cat = 0
if cat[node] == 0:
m_cat = 0
else:
m_cat = path_cat + cat[node]
# too more cat then leaves are not achievable
if m_cat > m:
return 0
isLeaf = True
sums = 0
# traverse edges belongs to node
for j in e[node]:
# ignore parent edge
if j == parent:
continue
# node has child
isLeaf = False
# dfs from child of node
sums += go_child(j, m_cat, node)
# achievable leaf
if isLeaf:
return 1
# return achievable leaves count
return sums
if __name__ == '__main__':
n, m = map(int, input().split())
cat = list(map(int, input().split()))
e = [[] for i in range(n)]
for i in range(n - 1):
v1, v2 = map(int, input().split())
# store undirected edge
e[v1 - 1].append(v2 - 1)
e[v2 - 1].append(v1 - 1)
# dfs from root
print(go_child(0, 0, -1))
So here is the thing, recursion in python is really bad. Default recursion depth is 1000, that means going deeper than 1000 calls will throw an exception. There are some ways to get around this
And as a final remark. You will have a much easier time coding in python if you submit with PyPy. Same code but it generally runs much faster. Also there are ways to speed up the IO. So as an example, this is what I would have used had I submitted your code: https://codeforces.com/contest/580/submission/55059855.
Btw avoid doing pop(0). It scales O(n), just with a small constant factor. Instead try to use deque for bfs.
Also, this time you don't even need bfs, you could just do dfs instead. So just switch pop(0) with pop() and your solution should run much faster.
You can also use a pointer, ptr = 0. Change bfs.pop(0) to bfs[ptr]. And at the end of your loop, ptr += 1. This will improve that problem.
Sorry for bad English.
If you want to do it that way then even this works:
The bad thing about writing it like this is that it is not obvious that it works (it does in python but probably not with an range based for loops in c++), but it is nice to not have to import deque and do popleft or mess around with indexing.