Codeforces problem python code optimization 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] &gt;= 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?
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 &gt;= 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 &gt;= 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))