codaility chlorum2014 challenge

chlorum2014.py

# https://codility.com/programmers/custom_challenge/chlorum2014
# not pass yet
	# SAMPLE  DATA
		# C=[1, 3, 0, 3, 2, 4, 4]
		# D=[6, 2, 7, 5, 6, 5, 2]
		# K=5

C=[1,3,0,3,2,4,4]
D=[6,2,7,5,6,5,2]
K=2

		
def solution(K, C, D):
	Top_D=sorted(D,reverse=True)[:K]
	Top_city_index_in_C = [idx for idx,val in enumerate(D) if val in Top_D]
	Visit=[x for x in Top_city_index_in_C for y in Top_city_index_in_C if x!=y and connected(C[x],C[y],C)]
   	return len(set(Visit))

def connected(P, Q ,C):
	return True if(C[P] == Q and P != Q) else False

if __name__ == '__main__':
	print 'you can visit : ' , solution(K,C,D)