2477 – Tricolor

Tanaka are un arbore (un tri) cu N noduri numerotate de la 1 la N. El vrea să coloreze nodurile arborelui în alb sau negru astfel încât numărul de perechi (neordonate) de noduri înfrățite să fie maxim. Două noduri sunt înfrățite dacă și numai dacă ambele sunt albe și fie sunt legate direct printr-o muchie, fie lanțul elementar unic dintre ele conține doar noduri negre.

Cerința

Dându-se un arbore cu N noduri, să se afle numărul maxim de perechi de noduri înfrățite ale sale care se poate obține.

Date de intrare

Fișierul de intrare tricolor.in va conține pe primul rând un număr natural nenul T ce reprezintă numărul de teste. Urmează T teste, fiecare test va descrie un arbore pentru care trebuie să se rezolve cerința. Pe primul rând al unui test apare un număr natural N ce reprezintă numărul de noduri ale arborelui din testul respectiv. Pe următoarele N-1 rânduri vor apărea câte o pereche de numere întregi x y separate printr-un spațiu, care indică existența unei muchii între nodul x și nodul y.

Date de ieșire

Fișierul de ieșire tricolor.out va conține T rânduri. Fiecare rând va conține soluția pentru câte un test, în aceeași ordine ca în fișierul de intrare.

Restricții și precizări

• 1 ≤ T ≤ 10
• 1 ≤ N ≤ 5 000
• Într-un test oarecare, 1 ≤ x, y ≤ N, x ≠ y
• Pentru 5 puncte, T = 1 și N ≤ 15
• Pentru alte 10 puncte, T = 1 și N ≤ 20
• Pentru alte 5 puncte, toți arborii descriși au exact 2 frunze și N ≤ 500
• Pentru alte 10 puncte, pentru toți arborii descriși există exact două noduri ale arborelui de care se leagă toate frunzele, situate la capetele unui lanț elementar și N ≤ 500.
• Pentru alte 50 de puncte, N ≤ 500
• Pentru alte 20 de puncte, nu există restricții suplimentare.

Exemplu:

tricolor.in

2
8
1 2
2 3
2 4
4 5
5 6
6 7
6 8
2
1 2

tricolor.out

7
1

Explicație

T = 2, avem două teste în fișierul de intrare. În primul test arborele are 8 noduri și cu o colorare optimă putem obține un număr maxim de 7 perechi de noduri înfrățite.

Perechile înfrățite sunt: (1, 3), (1, 4), (3, 4), (4, 5), (5, 7), (5, 8), (7, 8).

În al doilea test arborele are 2 noduri legate cu o singură muchie. Este optim să colorăm ambele noduri în alb obținând o pereche de noduri înfrățite.

Rezolvare

from collections import defaultdict

def maxFriendPairs(N, edges):
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    def dfs(node, parent):
        nonlocal ans
        black_count = 0
        white_count = 1

        for child in graph[node]:
            if child != parent:
                black_subtree_count, white_subtree_count = dfs(child, node)
                black_count += white_subtree_count
                white_count += black_subtree_count

        ans = max(ans, black_count * (black_count - 1) // 2 + white_count * (white_count - 1) // 2)
        return black_count, white_count

    ans = 0
    dfs(1, -1)
    return ans

def main():
    with open("tricolor.in", "r") as fin, open("tricolor.out", "w") as fout:
        T = int(fin.readline())
        for _ in range(T):
            N = int(fin.readline())
            edges = []
            for _ in range(N - 1):
                u, v = map(int, fin.readline().split())
                edges.append((u, v))
            fout.write(str(maxFriendPairs(N, edges)) + "\n")

if __name__ == "__main__":
    main()