Chương trình đếm số lượng nhóm động vật tối thiểu không có kẻ săn mồi bằng Python

Giả sử chúng ta có danh sách nums trong đó nums[i] cho biết kẻ săn mồi của động vật thứ i; nếu không có kẻ săn mồi thì giá trị là -1. Chúng ta cần tìm số lượng nhóm động vật nhỏ nhất sao cho không có động vật nào cùng nhóm với kẻ săn mồi trực tiếp hoặc gián tiếp của nó.

Ví dụ, nếu đầu vào là nums = [1, 2, -1, 4, 5, -1], thì đầu ra sẽ là 3, vì ta có thể chia thành các nhóm: [0, 3], [1, 4], [2, 5].

Quy trình giải quyết

1. Nếu danh sách A rỗng, trả về 0.
2. Xây dựng danh sách kề adj và tập gốc roots (những nút có giá trị -1).
3. Dùng DFS/BFS từ mỗi gốc để tìm độ sâu lớn nhất của cây.
4. Kết quả là độ sâu lớn nhất tìm được.

Ví dụ mã nguồn

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))  # Output: 3

Đầu vào

[1, 2, -1, 4, 5, -1]

Đầu ra

3