"The Story of a Tree" solved using depth-first search

Question

Found this problem in hackerrank.

One day Bob drew a tree, with \$n\$ nodes and \$n-1\$ edges on a piece of paper. He soon discovered that parent of a node depends on the root of the tree. The following images shows an example of that:

Learning the fact, Bob invented an exciting new game and decided to play it with Alice. The rules of the game is described below:

1. Bob picks a random node to be the tree's root and keeps the identity of the chosen node a secret from Alice. Each node has an equal probability of being picked as the root.

2. Alice then makes a list of \$g\$ guesses, where each guess is in the form \$u, v\$ and means Alice guesses that \${\rm parent}(v) = u\$ is true. It's guaranteed that an undirected edge connecting \$u\$ and \$v\$ exists in the tree.

3. For each correct guess, Alice earns one point. Alice wins the game if she earns at least \$k\$ points (i.e., at least \$k\$ of her guesses were true).

Alice and Bob play \$q\$ games. Given the tree, Alice's guesses, and the value of \$k\$ for each game, find the probability that Alice will win the game and print it on a new line as a reduced fraction in the format p/q.

Solution: There is a tree with some edges marked with arrows. For every vertex in a tree you have to count how many arrows point towards it. For one fixed vertex this may be done via one depth-first search (DFS). Every arrow that was traversed during DFS in direction opposite to its own adds 1. If you know the answer for vertex \$v\$, you can compute the answer for vertex \$u\$ adjacent to \$v\$ in \$O(1)\$.

It's almost the same as for \$v\$, but if there are arrows \$u→v\$ or \$v→u\$, their contributions are reversed. Now you can make the vertex \$u\$ crawl over the whole graph by moving to adjacent vertices in the second DFS.

def gcd(a, b):
    if not b:
        return a
    return gcd(b, a%b)

def dfs1(m, guess, root, seen):
    '''keep 1 node as root and calculate how many arrows are pointing towards it'''
    count = 0
    stack = []
    stack.append(root)
    seen.add(root)
    while len(stack):
        root = stack.pop(len(stack)-1)
        for i in m[root]:
            if i not in seen:
                seen.add(i)
                count += (1 if root in guess and i in guess[root] else 0)
                stack.append(i)
    return count

def dfs2(m, guess, root, seen, cost, k):
    '''now make every node as root and calculate how many nodes
       are pointed towards it; If u is the root node for which
       dfs1 calculated n (number of arrows pointed towards the root)
       then for v (adjacent node of u), it would be n-1 as v is the
       made the parent now in this step (only if there is a guess, if
       there is no guess then it would be not changed)'''
    stack = []
    stack.append((root, cost))
    seen.add(root)
    t_win = 0
    while len(stack):
        (root, cost) = stack.pop(len(stack)-1) 
        t_win += cost >= k
        for i in m[root]:
            if i not in seen:
                seen.add(i)
                stack.append((i, cost - (1 if root in guess and i in guess[root] else 0) + (1 if i in guess and root in guess[i] else 0)))
    return t_win


q = int(raw_input().strip())
for a0 in xrange(q):
    n = int(raw_input().strip())
    m = {}
    guess = {}
    seen = set()
    for a1 in xrange(n-1):
        u,v = raw_input().strip().split(' ')
        u,v = [int(u),int(v)]
        if u not in m:
            m[u] = []
        m[u].append(v)
        if v not in m:
            m[v] = []
        m[v].append(u)
    g,k = raw_input().strip().split(' ')
    g,k = [int(g),int(k)]
    for a1 in xrange(g):
        u,v = raw_input().strip().split(' ')
        u,v = [int(u),int(v)]
        if u not in guess:
            guess[u] = []
        guess[u].append(v)
    cost = dfs1(m, guess, 1, seen)
    seen = set()
    win = dfs2(m, guess, 1, seen, cost, k)
    g = gcd(win, n)
    print("{0}/{1}".format(win/g, n/g))

Answer 1

1. gcd is missing a docstring.

2. gcd is tail-recursive, but Python doesn't have tail-recursion elimination, and so:

   >>> gcd(3**7000, 2**10000)
   Traceback (most recent call last):
     File "<stdin>", line 1, in <module>
   [...]
   RecursionError: maximum recursion depth exceeded
   

   It is more robust to turn the tail-recursion into a loop:

   def gcd(a, b):
       """Return the greatest common divisor of a and b."""
       while b:
           a, b = b, a % b
       return a
   

   and then:

   >>> gcd(3**7000, 2**10000)
   1
   

3. Python has a built-in greater common divisor function, fractions.gcd, so you could use that and avoid writing your own.

4. The only use for gcd is to print a fraction in lowest terms. The built-in fractions.Fraction class can do that, so you could write:

   from fractions import Fraction
   print("{0.numerator}/{0.denominator}".format(Fraction(win, n)))
   

5. The names dfs1 and dfs2 are not very informative. Sure, they do some kind of depth-first search, but which one is which?

6. The docstrings for dfs1 and dfs2 don't explain the arguments m, guess, seen, cost, and k. What am I supposed to pass for these arguments?

7. The seen argument to dfs1 and dfs2 always gets an empty set which is not used again after the function returns. So it would better for these functions to use a local variable and initialize it with an empty set. Then there would be no risk of passing the wrong value for this argument.

8. Instead of creating an empty list and appending an item to it:

   stack = []
   stack.append(root)
   

   it's simpler to create a list with one item:

   stack = [root]
   

9. Instead of testing the length of a list:

   while len(stack):
   

   it's simpler to test the list itself:

   while stack:
   

10. Instead of computing the index of the last element of a list:

    root = stack.pop(len(stack)-1)
    

    you can use a negative index:

    root = stack.pop(-1)
    

    but by default the pop method pops the last element, so you can write:

    root = stack.pop()
    

11. Instead of writing this complex condition:

    count += (1 if root in guess and i in guess[root] else 0)
    

    you can take advantage of the fact that True equals 1 and False equals 0, and write:

    count += root in guess and i in guess[root]
    

12. The only use you make of guess is to test if it contains a directed edge. But this is slightly awkward because the test takes two steps. It would be better if you built a data structure that was better adapted to the use you intend to make of it. In this case, you could build a set of directed edges. So instead of:

    guess = {}
    for a1 in xrange(g):
        u,v = raw_input().strip().split(' ')
        u,v = [int(u),int(v)]
        if u not in guess:
            guess[u] = []
        guess[u].append(v)
    

    write:

    guess = set()
    for _ in range(g):
        u, v = map(int, raw_input().split())
        guess.add((u, v))
    

    and then in dfs1 you can write:

    count += (root, i) in guess
    

    (The construction of guess can be turned into a set comprehension, like this:

    guess = {tuple(map(int, raw_input().split())) for _ in range(g)}
    

    but I think the explicit loop is a bit clearer.)

13. Similarly, in dfs2 you have the complicated expression:

    stack.append((i, cost - (1 if root in guess and i in guess[root] else 0) + (1 if i in guess and root in guess[i] else 0)))
    

    Given the set of directed edges, this simplifies to:

    stack.append((i, cost - ((root, i) in guess) + ((i, root) in guess)))
    

14. The top-level code is hard to test, because you can't call it from a Python program or the interactive interpreter. Better to put it in its own function, conventionally called main.

15. The top-level code is full of one-letter variables (m, q, g, k, u, v). It is hard to guess what these mean. Although it makes sense to stick with the names from the HackerRank problem description, a comment or two would help remind the reader of what they mean.

16. When you don't use a loop variable like a1, it's conventional to use the name _.

17. Some of the repetition in the input code can be avoided by using collections.defaultdict:

    # Map from node to list of neighbours.
    tree = defaultdict(list)
    for _ in range(n - 1):
        u, v = map(int, raw_input().split())
        tree[u].append(v)
        tree[v].append(u)