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
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
3