Painting Fences - optimize for faster performance I am attempting this problem at Codeforces. I'm stuck as it gives a time limit exceeded error for large inputs though the logic I followed is similar to that of the editorial.

I'd be grateful for any pointers that can be useful in improving my logic and/or code performance.

import sys

count = 0

def count_and_split(p):
    h = min(p)
    n = len(p)
    global count
    if h&lt;n:
        count += h
    else:
        count += n
    p = [i-h for i in p] 
    e = []
    for i in p:
        if i == 0:
            k = p[:p.index(i)]
            if k:
                e.append(k)
            p = p[(p.index(i)+1):]
    if p:
        e.append(p)
    return e

def main():
    n = int(sys.stdin.readline())
    a = [map(int, sys.stdin.readline().split())]

    for part in a:       
        e = count_and_split(part)
        if count&gt;n:
            print n
            exit()
        for new_part in e:
            a.append(new_part)
    print count

if __name__ == '__main__':
    main()
Your logic seems badly organised. Let's try to improve this first :


Making things better organised

At the moment, when you get the input, you use them to define a and n such that, for example, a = [[2, 2, 1, 2, 1]] and n = 5. For the sake of clarity, let's try to isolate the different parts to make things easier to build. We could define a function getting the list [2, 2, 1, 2, 1] from the user and then write a function to handle this list.

Just doing this and moving your code around, here's what we get :

import sys

count = 0

def count_and_split(p):
    h = min(p)
    n = len(p)
    global count
    if h&lt;n:
        count += h
    else:
        count += n
    p = [i-h for i in p] 
    e = []
    for i in p:
        if i == 0:
            k = p[:p.index(i)]
            if k:
                e.append(k)
            p = p[(p.index(i)+1):]
    if p:
        e.append(p)
    return e

def get_nb_strokes(l):
    a = [l]
    n = len(l)
    global count
    count = 0
    for part in a:
        e = count_and_split(part)
        if count&gt;n:
            return n
        for new_part in e:
            a.append(new_part)
    return count

def get_list_from_user():
    n = int(sys.stdin.readline())
    return map(int, sys.stdin.readline().split())

def main():
    print(get_nb_strokes(get_list_from_user()))

if __name__ == '__main__':
    main()


It already looks much better doesn't it ? Well, it's not quite over yet.
Create tests :

Here's what your main function could look like as your change the code. You can run this quickly to ensure things seem to be working fine :

def main():
    assert get_nb_strokes([2, 2, 1, 2, 1]) == 3
    assert get_nb_strokes([2, 2]) == 2
    assert get_nb_strokes([1]) == 1

Add tests :

You can easily define new tests, it can also be a good idea to check how your code behaves on bigger inputs :

for i in range(2, 10):
    # fence of constant height
    for h in range(1,3):
        fence = [h] * i
        assert len(fence) == i
        assert get_nb_strokes(fence) == min(h,i)
    # fence of constant width
    for w in range(1,3):
        fence = [i] * w
        assert get_nb_strokes(fence) == min(w,i)
    # fence going up
    fence = range(1, i+1)
    assert get_nb_strokes(fence) == len(fence) == i
    # fence going down
    fence = range(i, 0, -1)
    assert get_nb_strokes(fence) == len(fence) == i
    # fence going up then down
    fence = range(1, i) + range(i, 0, -1)
    assert len(fence) == 2*i -1
    assert get_nb_strokes(fence) == i
    # fence going down then up
    fence = range(i, 0, -1) + range(2, i+1)
    assert get_nb_strokes(fence) == len(fence) == 2*i-1
    # L-shaped fence
    fence = [i] + [1]*i
    assert get_nb_strokes(fence) == 2
    # U-shaped fence
    fence = [i] + [1]*i + [i]
    assert get_nb_strokes(fence) == 3
    # W-shaped fence
    fence = [i] + [1]*i + [i] + [1]*i + [i]
    assert get_nb_strokes(fence) == 4

Don't use global variables

Having global variables makes things impossible to track and pretty hard to test. Let's try to improve this. It seems like the logic in count_and_split related to count could be done easily from the get_nb_strokes functions.

The updated code looks like :

def count_and_split(p):
    h = min(p)
    p = [i-h for i in p]
    e = []
    for i in p:
        if i == 0:
            k = p[:p.index(i)]
            if k:
                e.append(k)
            p = p[(p.index(i)+1):]
    if p:
        e.append(p)
    return e

def get_nb_strokes(l):
    a = [l]
    n = len(l)
    count = 0
    for part in a:
        count += min(min(part),len(part))
        e = count_and_split(part)
        if count&gt;n:
            return n
        for new_part in e:
            a.append(new_part)
    return count


Please note that at the moment, I still don't know how your code works, I've just changed the way it is organised.
A bit of cleanup

Now, we can move e = count_and_split(part) after the check about count as things are independent. We can also get rid of the temporary variable, use extend and rename the function to write :

def get_nb_strokes(l):
    a = [l]
    n = len(l)
    count = 0
    for part in a:
        count += min(min(part),len(part))
        if count&gt;n:
            return n
        a.extend(split(part))
    return count



I'll try to have another loop ASAP.