Codeforces problem python code optimization

Question

I have solved one problem on Codeforces and I guess the output is quite correct. However, it is slow for input in range of thousands. Since, I am learning Python now, it would be great help if one could help me optimise this code. Any notes, tips or suggestions on the sidelines are always welcome.

The question is a practice question and hence does not violate any contest rules.

import sys
import time

_no_of_ants, _no_of_queries = (input().split())
no_of_ants = int(_no_of_ants)
no_of_queries = int(_no_of_queries)

class ants_cord():
    def __init__(self,point,no_of_ants):
        self.co_ant = {point:no_of_ants}

    def ants_at_point(self,point):
        return self.co_ant[point]

    def move_ants(self,point):
        x1,y1 = point
        if (x1+1,y1) in self.co_ant.keys():
            self.co_ant[((x1+1),y1)] += 1
        else:
            self.co_ant[((x1+1),y1)] = 1
        if (x1,y1+1) in self.co_ant.keys():
            self.co_ant[((x1),y1+1)] += 1
        else:
            self.co_ant[((x1),y1+1)] = 1
        if (x1,y1-1) in self.co_ant.keys():
            self.co_ant[((x1),y1-1)] += 1
        else:
            self.co_ant[((x1),y1-1)] = 1
        if (x1-1,y1) in self.co_ant.keys():
            self.co_ant[((x1-1),y1)] += 1
        else:
            self.co_ant[((x1-1),y1)] = 1
        self.co_ant[point] -= 4

    def check_status(self):
        for key in self.co_ant.keys():
            if k.co_ant[key] >= 4:
                return key
        else: return False
start = time.time()
k = ants_cord((0,0),no_of_ants)
point = k.check_status()
while point: 
    k.move_ants(point)
    point = k.check_status()
end = time.time()
print("Sec:{}".format(end-start))
while no_of_queries:
    _x,_y = input().split()
    x = int(_x)
    y = int(_y)
    if (x,y) in k.co_ant:
        print(k.co_ant[(x,y)])
    else:
        print('0')
    no_of_queries -= 1

EDIT: I updated the code according to suggestion. But I still get a TLE error in test case 4. Should I employ other algorithm then?

Answer 1

First answer was about micro-optimisation.

Here are a few improvements you could perform on the algorithm itself :

• You could simulate many iterations in one go.

In order to do so,

def move_ants(g,point):
    x,y = point
    for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
        g[a,b] = 1 + g.get((a,b),0)
    g[point] -= 4

could be changed for :

def move_ants(g,point):
    x,y = point
    nb_ants = g[point]
    nb_ants_per_dir = nb_ants/4
    for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
        g[a,b] = nb_ants_per_dir + g.get((a,b),0)
    g[point] = nb_ants - 4 * nb_ants_per_dir

• check_status could return a set of points to be handled

You'd just need to do something like :

def check_status(g):
    return [k for k,v in g.iteritems() if v >= 4]

• Also, you might want to keep track of the cells you've updated in order no to loop over the all dict but only the candidates. This would need quite a lot of rewritting.

• Also, you could take into account the fact that the problem offer many axis of symetry : obviously the x and y axis but also the y=x and the y=-x. Thus you could restrict yourself to 1/8 of the plan corresponding to 0 < x < y

Edit

After testing, this code on 5555 ants:

  def move_ants(g,points):
      cand=set()
      for point in points:
          x,y = point
          nb_ants = g[point]
          nb_ants_per_dir = nb_ants/4
          for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
              new_nb = g.get((a,b),0) + nb_ants_per_dir
              g[a,b] = new_nb
              if new_nb >= 4:
                  cand.add((a,b))
          g[point] = nb_ants - 4*nb_ants_per_dir
      return cand

  g = {(0,0) : no_of_ants} # or you could use default dict
  points = [(0,0)]
  while points:
      points = move_ants(g,points)
  print("Sec:{}".format(time.time()-start))
  for q in xrange(no_of_queries):
      x,y=1,1
      print(g.get((x,y),0))

appears to be 21 times faster than the solution provided in my other answer and 141 times faster than yours.

Answer 2

A few things:

• You should probably put the start = time.time() as early as possible if you want it to make sense.

• As your class consists on a simple dict, there is no or little value in defining it. (In a "real life" scenario, it would be considered as a good practice).

Things got 25% quicker on average starting with 1000 ants.

• You could make your code more concise by using structure to isolate the parts that change.

move_ants() would look like this :

def move_ants(g,point):
    x,y = point
    for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
        if (a,b) in g.keys():
            g[a,b] += 1
        else:
            g[a,b] = 1
    g[point] -= 4

(Not performance impact)

• In check_status(), instead of iterating on keys and then retrieve the corresponding value, you could use iteritems() :

check_status() would be like this :

def check_status(g):
    for k,v in g.iteritems():
        if v >= 4:
            return k
    else: return False

(This does improve performances)

• To get the values from your dict with default values, you could use get(). You'd have something like print(g.get((x,y),0)) and g[a,b] = 1 + g.get((a,b),0). (I really didn't expect so but this made the script more than twice faster).

At the end, my script is like this (I've changed other details but it's mostly personal preferences) :

#!/usr/bin/python

import sys
import time

start = time.time()

##_no_of_ants, _no_of_queries = (input().split())
##no_of_ants = int(_no_of_ants)
##no_of_queries = int(_no_of_queries)

no_of_ants = 1000
no_of_queries = 1

def move_ants(g,point):
    x,y = point
    for a,b in [(x+1,y),(x,y+1),(x,y-1),(x-1,y)]: # or you could define x generxtor
        g[a,b] = 1 + g.get((a,b),0)
    g[point] -= 4

def check_status(g):
    for k,v in g.iteritems():
        if v >= 4:
            return k
    else: return False

g = {(0,0) : no_of_ants} # or you could use default dict
point = check_status(g)
while point:
    move_ants(g,point)
    point = check_status(g)
print("Sec:{}".format(time.time()-start))
for q in xrange(no_of_queries):
    ###_x,_y = input().split()
    ###x = int(_x)
    ###y = int(_y)
    x,y=1,1
    print(g.get((x,y),0))