"Tractor" Python implementation

Question

This is the Tractor problem from the Kattis problems archive.

Problem Statement

Bessie the Cow has stolen Farmer John’s tractor and is running wild on the coordinate plane! She, however, is a terrible driver, and can only move according to the following rules:

Each of her movements is in the same direction as either the positive x-axis or the positive y-axis. Her \$\text{nth}\$ movement takes her \$2^{n−1}\$ units forward in her chosen direction. (On her first movement, \$n = 1\$, so she moves 1 unit.) Farmer John’s farm is on the coordinate plane, in the shape of a rectangle with corners at \$(0, 0), (A, 0), (0, B) \text{and} (A, B)\$. If Bessie starts at (0, 0), how many points inside the farm, including the boundary, could she reach?

Input Format

The input begins with an integer \$N (1 \le N \le 100)\$ on a line by itself, indicating the number of test cases that follow. Each of the following N lines contains two space separated integers A and B \$(1 \le A, B \le 10^8)\$, describing the upper-right corner of Farmer John’s farm.

Output Format

Output N lines, with the Nth line containing the number of points that Bessie could possibly reach in the Nth test case.

In the first test case of the sample, Bessie can reach the following six points: \$(0, 0), (0, 1), (1, 0), (1, 2), (2, 1) and (3, 0)\$.

Sample Input

2
2 3
7 7

Sample Output

6
15

Solution

def next_steps():
    n = 1
    while True:
        yield n
        n <<= 1

def moves(a, b):
    current_positions = [(0, 0)]
    moves = []
    for step in next_steps():
        next_positions = []
        while current_positions:
            x, y = current_positions.pop(0)
            moves.append((x, y))
            if x + step <= a:
                next_positions.append((x+step, y))
            if y + step <= b:
                next_positions.append((x, y+step))
        if not next_positions:
            break
        else:
            current_positions = next_positions
    return moves

if __name__ == '__main__':
    for test_case in xrange(int(raw_input())):
        a, b = map(int, raw_input().strip().split())
        print len(moves(a, b))

It's working fine on smaller input but gets timed out on very large instances.

Answer 1

for loop

Use a for loop when iterating over a collection:

    for (x, y) in current_positions:
        moves.append((x, y))

for has less power than while (theoretically) and its simplicity makes it a better choice over while.

Avoid premature optimization

Just multiply by two:

n *= 2

Some tests:

>>> import timeit
>>> timeit.timeit("10 << 1")
0.029866752000089036
>>> timeit.timeit("10 * 2")
0.03209113599950797
>>> 0.029866752000089036 / 0.03209113599950797
0.9306854079751793

<< is just 7% faster than *2

And the majority of your time is spent branching if x + step <= a: and playing with lists moves.append((x, y)), so the final improvement would be tiny, not worth the loss in readability.

Optimizing the algorithm

Your algorithm has exponential time complexity:

print(step, len(current_positions))

Gives:

1 2
2 4
4 8
8 16
16 32
32 64
64 74
128 0

At each iteration, the number of current_positions is doubling, and this can get out of hand really quickly.

Arguments name

a and b are not acceptable argument names. I suggest width and height.

Some plotting always gives ideas

Following @ErikR suggestion, I made a graph of the reachable points and noticed a serious regularity, still doing graphs is a nice skill to have, so I encourage you to reproduce this graph by yourself.

Answer 2

Note that for \$a = b = 10^6\$ the answer is 2000001 which is also the size of your moves list. Since you are only interested in its length, just use the variable to store the length of the list - i.e. replace moves.append(...) with moves += 1. This will enable you to solve the problem for larger values of a and b.

But to solve the problem for really large values you need a better algorithm. One way to see this is to note that moves can increase only by 1 on each iteration of the loop. That means for \$a = b = 10^6\$ your problem will be performing at least 2 million loop iterations. On my machine 200K iterations already takes 8 seconds, so that shows that counting by 1 isn't a feasible way to solve this problem.

For a hint to a better method, plot out the points for a = 10, b = 10 and see if you can discern a pattern.

Answer 3

A better algorithm

Thanks to ErikR and Caridorc for providing the base and image used to produce this answer:

From image one sees a clear diagonal pattern which can be expressed using the following function based upon \$x\$ and \$n\$:

$$ f(x, n) = 2^n -1 - x $$

Added: This function denotes the \$y\$ coordinate of the dots in the image. Let \$x\$ range from \$0\$ through \$A\$, first for \$n=0\$ which gives the \$(0, 0)\$ point (and a lot of points outside the coordinate plane). For \$n=1\$ it gives the \$(1, 0)\$ and \$(0, 1)\$, and so it continues. To calculate for \$f(0, 3)\$ you need to find to use \$2^3\$ (aka print 2**3 in python), which then gives: $$ f(0, 3) = 2^n - 1 - x = 8 -1 -0 = 7 $$

Which in turn gives us the \$y\$ of the fourth diagonal in the diagram (counting diagonals from the lower left including the one consisting of just the point \$(0, 0)\$, that is the point \$(0, f(0, 3)) = (0, 7)\$

Further since this function has a constant drop of 1 pr increase in \$x\$, the difference between the \$x_{start}\$ and the \$x_{end}\$ plus 1 of any given diagonal indicates how many places/moves this diagonal contributes to the final amount of legal moves. Finally, the \$x_{start}\$ needs either to be on the left or top edge or it is uninteresting, and the \$x_{end}\$ needs to be at the bottom or right edge to have any interest.

Added: \$x_{start}\$ is the x coordinate of the leftmost point in any given diagonal, and the \$x_{end}\$ is the x coordinate of the rightmost in the same diagonal. For \$n\$ diagonals you'll have \$n\$ pairs of \$x_{start}\$ and \$x_{end}\$, in the image \$5\$ diagonals are shown including the \$(0, 0)\$ diagonal.

Formalisation of previous statements:

• Left edge – If \$f(0, n) < B\$ for current \$n\$, then \$x_{start}\$ is 0
• Top edge – If \$f(0, n) > B\$ for current \$n\$, then find the \$x\$ where \$f(x, n) == B\$. That is let \$x_{start} = 2^n - 1 - B\$
• If \$x_{start} > A\$, then it is uninteresting as the diagonal doesn't cross our coordinate plane, we can then use \$x_{start} = A + 1\$
• Bottom edge – Find \$f(x, n) = 0\$ for current \$n\$. That is \$x_{end}\ = 2^n - 1\$
• If \$x_{end} > A\$ then limit it to \$x_{end} = A\$
• Places contribution is only valid if \$x_{start} <= A\$ and \$x_{start} <= x_{end}\$, and if valid the contribution is: \$x_{end} - x_{start} + 1\$

Do the calcuation for \$n = 3\$

• Left edge – \$f(0, 3) = 2^3 - 1 - 0 = 7\$, which is lower than B, let \$x_{start} = 0\$
• Bottom edge – \$x_{end} = 2^3 -1 = 7\$, which is lower than A, let \$x_{end} = 7\$
• Valid contribution – \$7 - 0 + 1 = 8\$

Do the calculation for \$n = 4\$

• Left edge – \$f(0, 4) = 2^4 - 1 - 0 = 15\$, which is above B
• Top edge – Let \$x_{start} = 2^4 - 1 - 10 = 5\$, which is lower than \$A\$ so it's valid
• Bottom edge – Let \$x_{end} = 2^4 - 1 = 15\$, which is larger than \$A\$, so limit it to \$x_{end} = 10\$
• Valid contribution – \$10 - 5 + 1 = 6\$

Code refactor

When refactoring for coding an optimal calculation I add the following statements:

• When the \$x_{start}\$ leaves the left edge, it will never come back down again, and similarily when it passes \$A\$ it is for ever lost
• When the \$x_{end}\$ crosses the right edge, it can always be limited to \$A\$

This leads to the following code using generators for calculating the next diagonal x pairs:

def x_start_generator(a, b):
    """Generate start x-coordinate of the n'th diagonal"""

    above_b = False    # Indicates if an earlier x_start=0 has y > b
    beyond_a = False   # Indicates if an earlier x_start > a
    two_pow_n = 1      # Iterates 2**n, instead of calculating it

    while True:
        # print('two_pow_n: {}, diag: {}, above_b: {}, beyond_a: {}'.format(
        #      two_pow_n, diagonal_function(0, two_pow_n), above_b, beyond_a))

        # if x_start has been bigger than a, then it will continue to be bigger
        # so we always return a+1
        if beyond_a:
            yield a+1
            continue

        # Verify if x_start is on left edge (that is 'not above_b')
        if not above_b:
            above_b =  (two_pow_n - 1) > b


        # If on left edge, return start_x = 0 ...
        if not above_b:
            yield 0
        else:

            # ... else find x_start on line at height b
            x_start = two_pow_n - 1 - b

            # Check if x_start is beyond the value of a, and store this
            beyond_a = x_start > a

            # If not beyond return x_start matching the top edge
            if not beyond_a:
                yield x_start

        # Prepare for next n
        two_pow_n *= 2


def x_end_generator(a, b):
    """Generate end x-coordinate of the n'th diagonal"""

    beyond_a = False  # Indicates if an earlier x_end > b
    two_pow_n = 1

    while True:
        # If x_end has once been larger than a, limit it to a
        if beyond_a:
            yield a
            continue

        # Find end coordinate when diagonal crosses y=0
        x_end = two_pow_n - 1

        # Check if we've passed a, and store this
        beyond_a = x_end > a

        # If not passed, return x coordinate
        if not beyond_a:
            yield x_end

        # Prepare for next iteration
        two_pow_n *= 2


def find_legal_places(a, b, print_xtras=False):
    """Return number of legal places within (0, 0) -> (a, b)

    Using the x_start_generator() and x_end_generator() find the start
    and end coordinators of the n'th diagonal. If it's a valid diagonal
    add the legal places from this diagonal to the total. When not any
    more valid coordinates are found, return the total.

    If wanted, print_xtras=True, you can get print outs of the coordinates
    as we tag along, and a nice total.
    """
    x_start_range = x_start_generator(a, b)
    x_end_range = x_end_generator(a, b)

    legal_places = 0
    n = 0

    while True:
        x_start = next(x_start_range)
        x_end = next(x_end_range)

        if x_start < a and x_start <= x_end:
            legal_places += x_end - x_start + 1
            if print_xtras:
                print('n: {: 4d}, x_start: {: 12d}, x_end: {: 12d}, legal places: {: 12d}'.format(
                      n, x_start, x_end, legal_places))
        else:
            break

        n += 1
    if print_xtras:
        print('    For A={} and B={} there are {} legal places\n'.format(a, b, legal_places))

    return legal_places

def main():
    print 'Sample Output 1'
    for (a, b) in [(2, 3), (7, 7)]:
       print(find_legal_places(a, b))

    print

    find_legal_places(10, 10, True)

    print('a=10**1000, b=10**1000, legal_places={}'.format(find_legal_places(10**1000, 10**1000)))


if __name__ == '__main__':
    main()

When running my main() with a few test cases of varying magnitude, it all completed in less then a second. Cases run are (a, b) in [(2, 3), (7, 7), (10, 10), (10**1000, 10**1000). The latter would according to ErikR require quite a lot of iterations...

The output of this run was as follows:

Sample Output 1
6
15

n:    0, x_start:            0, x_end:            0, legal places:            1
n:    1, x_start:            0, x_end:            1, legal places:            3
n:    2, x_start:            0, x_end:            3, legal places:            7
n:    3, x_start:            0, x_end:            7, legal places:           15
n:    4, x_start:            5, x_end:           10, legal places:           21
    For A=10 and B=10 there are 21 legal places

a=10**1000, b=10**1000, legal_places=20000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001

That last number is 2 followed by 999 0's and ending with a 1, or put another way 2 * 10^1000 + 1. It seems like when \$a = b = x\$, then there is \$2 * x + 1\$ legal places¹.

_{¹I tested for all cases up through \$10^{1000}\$}

Answer 4

Style and Code review

Here are some comments on how to make your current logic somewhat faster, even though not in the same leagues as you get when switching to a better algoritm. This is to help you the next time you find a similar problem.

• General style is good – You might get a few comments on using variable names like a, b, ... but as they refer directly to your problem statement I find it acceptable
• Include problem statement in docstrings – In the original file of my answer I've included the problem statement at the start of the file in a docstring, this allows me when I later revisit the file to understand what problem I was trying to solve. Even just adding a link to the problem is better than nothing, but in general I would suggest making a textual copy in docstrings
• Add docstrings to methods – Always document what it does and returns if is a slightly complex method
• Avoid code in the __main__ block – Move code here to a well named function, and it is easier to understand what is happening. The map(int, raw_input().strip().split()) is close to garbage when skimming it
• Don't store any more data then needed – This is the major performance issues in your code. You store every legal move, when you only need the count of all moves. Storing current_positions and next_positions is good, storing moves not good at all
• Use appropriate loop methods – Use for position in current_positions instead of a while loop. It is more intuitive, and that is what you are actually doing. This also skips the unnatural pop which modifies the list on every iteration, which is unneccessary

You should try refactoring your code according to these suggestions, and you'll see a massive increase in performance when going to the larger versions.