m-aryTreeQues.py

class Node:
  def __init__(self, parent, children, country):
    self.parent = parent
    self.children = children
    self.country = country
    self.isLocked = 0
    self.ldc = 0
  
  def lock(self, user_id):
    if self.isLocked != 0:
      return False

    if self.ldc > 0:
      return False
    
    ancestor = self.parent
    while ancestor:
      if ancestor.isLocked:
        return False
      ancestor = ancestor.parent
    
    ancestor = self.parent
    while ancestor:
      ancestor.ldc += 1
      ancestor = ancestor.parent
    
    self.isLocked = user_id
    return True



  def unlock(self, user_id):
    if self.isLocked == 0 and self.isLocked != user_id:
      return False

    ancestor = self.parent

    while ancestor:
      ancestor.ldc -= 1
      ancestor = ancestor.parent

    self.isLocked = 0
    return True

  def upgradeLock(self, user_id):
  
    def traverse(root):
      if root:
        for c in root.children:
          if mapping.get(c, None):
            if mapping[c].isLocked != user_id:
              return False
          traverse(mapping.get(c, None))
      return True

    isValid = traverse(mapping.get(self.country, None))

    if not isValid:
      return False

    ancestor = self.parent
    while ancestor:
      ancestor.ldc += 1
      ancestor = ancestor.parent
    
    self.isLocked = user_id
    return True


N = int(input())
m = int(input())
queries_total = int(input())

nodes_arr = []
mapping = {}

for i in range(0, N):
  nodes_arr.append(input())

queries = []

for i in range(0, queries_total):
  queries.append(list(input().split(" ")))


for idx in range(0, N):
  cntr = nodes_arr[idx]
  parent = None
  children = [None for _ in range(m)]
  for x in range(m):
    val = m*idx + x + 1
    if val < N:
      children[x] = nodes_arr[val]

  pv = idx//m 
  parent = mapping.get(nodes_arr[pv], None)

  mapping[cntr] = Node(parent, children, cntr)
  
for query in queries:
  cntr = query[1]
  if int(query[0]) == 1:
    print("True" if mapping[cntr].lock(query[2]) else "False")
  elif int(query[0]) == 2:
    print("True" if mapping[cntr].unlock(query[2]) else "False")
  elif int(query[0]) == 3:
    print("True" if mapping[cntr].upgradeLock(query[2]) else "False")

# Input  
# 7
# 2
# 5
# World
# Asia
# Africa
# China
# India
# SouthAfrica
# Egypt
# 1 China 9
# 1 India 9
# 3 Asia 9
# 2 India 9
# 2 Asia 9

# Output
# true
# true
# true
# false
# true