Program to count minimum number of animals which have no predator in Python

Python Server Side Programming Programming

Suppose we have a list of numbers called nums where nums[i] shows the predator of the ith animal and if there is no predator, it will hold −1. We have to find the smallest number of groups of animals such that no animal is in the same group with its direct or indirect predator.

So, if the input is like nums = [1, 2, −1, 4, 5, −1], then the output will be 3, as we can have the groups like: [0, 3], [1, 4], [2, 5].

To solve this, we will follow these steps −

if A is empty, then
- return 0
adj := a blank map
vis := a new set
roots := a new list
for each index i and value a in A, do
- if a is same as −1, then
  - insert i at the end of roots
- insert a at the end of adj[i]
- insert i at the end of adj[a]
best := −infinity
for each root in roots, do
- stk := a stack and insert [root, 1] into it
- while stk is not empty, do
  - (node, d) := popped element of stk
  - if node is in vis or node is same as −1, then
    - come out from the loop
  - best := maximum of best and d
  - insert node into vis
  - for each u in adj[node], do
    - push (u, d + 1) into stk
return best

Let us see the following implementation to get better understanding −

Example

Live Demo

from collections import defaultdict
class Solution:
   def solve(self, A):
      if not A:
         return 0
      adj = defaultdict(list)
      vis = set()
      roots = []
      for i, a in enumerate(A):
         if a == -1:
            roots.append(i)
         adj[i].append(a)
         adj[a].append(i)
      best = −float("inf")
      for root in roots:
      stk = [(root, 1)]
      while stk:
         node, d = stk.pop()
         if node in vis or node == −1:
            continue
         best = max(best, d)
         vis.add(node)
         for u in adj[node]:
            stk.append((u, d + 1))
   return best
ob = Solution()
nums = [1, 2, −1, 4, 5, −1]
print(ob.solve(nums))

Input

[1, 2, −1, 4, 5, −1]

Output

Arnab Chakraborty

Updated on: 21-Oct-2020

604 Views

Kickstart Your Career

Get certified by completing the course

Get Started