Implementing a category tree

I am implementing a continent-country relation category map In my implementation a method add_country(), and as a parameter of add_country I provide the name of the country and continent it belongs to, for example, add_country(Uganda, Africa). This adds a node Africa->Uganda.

If I pass a similar function, with the second parameter as 0, ex. add_country(Asia, 0), it should add the first parameter as a category head or parent. The get_country() method should return the name of all countries belonging to that continent (or children).

So far I have made this:

class Node:
    def __init__(self, val):
        self.children = []
        self.value = val

    def add_child(self, val):
        n = Node(val)
        self.children.append(n)

    def print_stat(self):
      print(self.children)
      print(self.value)

class CountryTree:

    def __init__(self):
        self.roots = []


    def add_country(self, country, parent):
      if parent == None:
        root = Node(country)
        self.roots.append(root)
      else:
        par = Node(parent)
        par.add_child(country)



    def get_country(self, parent):
      par = Node(parent)
      return par.children



map = CountryTree()
map.add_country('Asia', None)
map.add_country('China', 'Asia')
map.add_country('Russia', 'Asia')
map.get_country('Asia') #Should return China and Russia

But running this is returning me with an empty array.

Solution

I think the following does what you want. Most of the significant changes were to the CountryTree class. I changed the list of continents you had into a dictionary. This makes determining if one is present easy and fast — something your add_country() method wasn't doing — which was one of the major things wrong with it. I also changed its name from self.roots to self.continents to be more descriptive. I also gave the parent argument a default value of None which makes providing it when adding continent nodes optional.

Another important change was to make the code follow the PEP 8 - Style Guide for Python Code guidelines to make it more readable and understandable. I strongly suggest that you read and follow them yourself.

class Node:
    def __init__(self, value):
        self.children = []
        self.value = value

    def add_child(self, value):
        self.children.append(Node(value))

    def __repr__(self):
        classname = type(self).__name__
        return (f'{classname}({self.value!r}, {self.children})' if self.children else
                f'{classname}({self.value!r})')

    def print_stat(self):
        print(self.children)
        print(self.value)


class CountryTree:
    def __init__(self):
        self.continents = {}  # Dictionary mapping continents to countries.

    def add_country(self, country, parent=None):
        if not parent:  # Adding a continent?
            continent = country
            self.continents[continent] = Node(continent)
        else:  # Adding a country to a continent.
            continent = self.continents.get(parent, None)
            if not continent:
                raise KeyError(f'No continent named {parent} exists')
            continent.add_child(country)

    def get_country(self, parent):
        continent = self.continents.get(parent, None)
        if not continent:
            raise KeyError(f'No continent named {parent} exists')
        return continent.children


map = CountryTree()
map.add_country('Asia')
map.add_country('China', 'Asia')
map.add_country('Russia', 'Asia')
result = map.get_country('Asia') # Should return China and Russia
print(result)

Output:

[Node('China'), Node('Russia')]