CODEWITHSUNDEEP
python
Sajad Rastegar
Sajad Rastegar

Reputation: 3154

A better way to write this program

I have converted the following block of codes from a java program. How can I write the name of the countries in Map using their names, instead of their IDs?

from collections import defaultdict
colors = ['Red', 'Yellow', 'Green', 'Blue']
mapColors = defaultdict(str)

def okToColor(Map ,country, color):
    for c in Map[country]:
        if mapColors[c] == color: return False
    return True

def explore(Map, country, color):
    if country >= len(Map): return True
    if okToColor(Map, country, color):
        mapColors[country] = color
        for color in colors:
            if explore(Map, country + 1, color): return True
    return False

def printMap():
    for c in mapColors:
        print c, mapColors[c]

Map = [[1, 4, 2, 5], [0, 4, 6, 5], [0, 4, 3, 6, 5], [2, 4, 6],
        [0, 1, 6, 3, 2], [2, 6, 1, 0], [2, 3, 4, 1, 5]]
result = explore(Map, 0, 'Red')
print result
printMap()

I want't the Map to be a graph like this:

Map = { 'A':['B', 'C'], 'B':['A','D'], ...}

where A, B, C, D are the name of the countries.

Upvotes: 1

Views: 569

Answers (1)

unutbu
unutbu

Reputation: 880807

The main idea is to define a mapping between countries and numerical indices:

countries = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
cindex = dict(zip(countries, range(len(countries))))
# {'A': 0, 'B': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6}

Then, with only little changes, you can use your original code. Where before country was a numerical index, now you put cindex[country] when you need the numerical index.

And when you need to reverse the mapping, countries[index] gives you the string name of the country.

from collections import defaultdict

colors = ['Red', 'Yellow', 'Green', 'Blue']
mapColors = defaultdict(str)


def okToColor(Map, country, color):
    for c in Map[country]:
        if mapColors[c] == color:
            return False
    return True


def explore(Map, country, color):
    if cindex[country] >= len(Map):
        return True
    if okToColor(Map, country, color):
        mapColors[country] = color
        for color in colors:
            try:
                next_country = countries[cindex[country] + 1]
            except IndexError:
                return True
            if explore(Map, next_country, color):
                return True
    return False


def printMap():
    for c in mapColors:
        print c, mapColors[c]

countries = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
cindex = dict(zip(countries, range(len(countries))))
Map = [[1, 4, 2, 5], [0, 4, 6, 5], [0, 4, 3, 6, 5], [2, 4, 6],
      [0, 1, 6, 3, 2], [2, 6, 1, 0], [2, 3, 4, 1, 5]]
Map = {countries[idx]: [countries[n] for n in item] for idx,
       item in enumerate(Map)}
result = explore(Map, 'A', 'Red')
print(result)
printMap()

yields

True
A Red
C Yellow
B Yellow
E Green
D Red
G Blue
F Green

Upvotes: 2

Related Questions