Painting Fences - optimize for faster performance

Question

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<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>n:
            print n
            exit()
        for new_part in e:
            a.append(new_part)
    print count

if __name__ == '__main__':
    main()

Answer 1

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<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>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>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>n:
              return n
          a.extend(split(part))
      return count
  

I'll try to have another loop ASAP.