CF 102420B - Сильная группа

CF 102420B - \u0421\u0438\u043b\u044c\u043d\u0430\u044f \u0433\u0440\u0443\u043f\u043f\u0430

Rating: -
Tags: -
Solve time: 10m 9s
Verified: yes

Solution

Problem Understanding

У нас есть дерево из n комнат. В каждой комнате находится один эльф с силой w[i]. Нужно выбрать хотя бы две комнаты так, чтобы выбранные комнаты образовывали связное поддерево. Среди всех таких групп требуется найти максимальное среднее значение сил.

Связность здесь существенно ограничивает выбор. Если выбраны две комнаты, то все вершины на пути между ними тоже обязаны быть выбраны. При этом выбранная группа может состоять из двух соседних вершин, из длинного пути, из вершины с несколькими ветвями или вообще из произвольного связного поддерева.

n достигает 200000, поэтому перебор всех групп невозможен. Уже на простой цепочке число связных групп размера хотя бы два равно

[ \frac{n(n-1)}2. ]

При n = 200000 это 19 999 900 000 вариантов. Для произвольного дерева количество связных подмножеств вообще может быть экспоненциальным. Значит, алгоритм должен работать почти линейно на каждой проверке, а само число проверок должно быть небольшим.

Силы лежат от 0 до 10^9, поэтому ответ также лежит в этом диапазоне. Требуемая точность 10^-6 естественно подсказывает двоичный поиск по вещественному ответу. При примерно 60 итерациях интервал длиной 10^9 становится меньше 10^-9, чего с большим запасом хватает для требуемой точности.

Есть несколько случаев, в которых наивная логика легко ломается. Во-первых, нельзя разрешить группу из одной вершины. Например,

2
100 0
1 2

Единственная допустимая группа состоит из двух эльфов, поэтому ответ равен 50. Если при проверке некоторого значения просто искать связное множество с положительной преобразованной суммой, одна вершина веса 100 может ошибочно объявить значение 75 достижимым.

Во-вторых, нельзя просто взять две самые сильные вершины. Например,

Две вершины силы 10 не соединены напрямую, поэтому их нельзя выбрать без промежуточных вершин. Максимальное допустимое среднее равно 5, например для группы {1,2} или {3,4}. Ответ 10 был бы ошибкой.

Наконец, оптимальная группа не обязана состоять из двух вершин. Например,

Любая пара соседних вершин имеет среднее 5, но группа из всех четырех вершин имеет среднее 7.5. Значит, проверка только рёбер тоже недостаточна.

Approaches

Самый прямой способ заключается в переборе всех связных групп, вычислении суммы их сил и делении на размер группы. На цепочке число кандидатов уже равно n(n-1)/2, то есть при n = 200000 примерно 2 * 10^10. Если для каждой группы ещё отдельно искать сумму, работа становится ещё больше. Для общего дерева число связных групп может быть экспоненциальным, поэтому такой подход не просто медленный, он принципиально не масштабируется.

Нужно убрать среднее из задачи. Предположим, что мы хотим проверить, существует ли группа со средним не меньше некоторого числа x. Для выбранной группы S имеем

[ \frac{\sum_{v\in S} w_v}{|S|} > x ]

тогда и только тогда, когда

[ \sum_{v\in S}(w_v-x)>0. ]

Значит, для фиксированного x можно заменить силу каждой вершины на

[ a_v=w_v-x ]

и искать связное множество размера хотя бы два с положительной суммой.

Теперь задача стала обычной задачей о максимальной сумме связного поддерева, но с дополнительным условием размера. На дереве такую задачу решает динамика снизу вверх.

Зафиксируем корень дерева. Пусть dp[u] означает максимальную сумму связного множества, которое содержит u и целиком лежит в поддереве u. Если мы уже посчитали значение для сына v, то его можно присоединить к u, только если dp[v] положительно. Если оно отрицательно, добавление этой ветки только ухудшит сумму.

Получаем

[ dp[u]=a_u+\sum_{v\text{ child of }u}\max(0,dp[v]). ]

Но здесь появляется тонкость. dp[u] может соответствовать множеству из одной вершины. Для проверки исходной задачи этого недостаточно. Поэтому отдельно рассматриваем лучший вариант, который обязательно присоединяет хотя бы одного сына.

Если у u есть сын с неотрицательным dp, то к u можно присоединить все положительные сыновья и хотя бы одного сына с нулевым значением. Если все значения отрицательны, никакого допустимого множества с положительной суммой через u получить нельзя.

Таким образом, одна проверка выполняется за O(n). Само условие достижимости монотонно: если существует группа со средним больше x, то она существует и для любого меньшего значения. Это позволяет использовать двоичный поиск.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	Exponential in general, at least `Θ(n²)` even on a path	`O(n)`	Too slow
Binary Search + Tree DP	`O(n log(10^9 / eps))`	`O(n)`	Accepted

Algorithm Walkthrough

Корнем дерева выбираем вершину 1 и один раз строим массив parent и порядок обхода order. Используем итеративный обход, потому что при n = 200000 рекурсивный DFS в Python может упереться в ограничение глубины рекурсии.
Для текущего кандидата x заменяем мысленно вес каждой вершины на w[u] - x. Наша задача теперь состоит в том, чтобы найти связную группу из хотя бы двух вершин с суммой больше нуля.
Обрабатываем вершины в обратном порядке order, то есть сначала листья, затем их родителей. Для каждой вершины вычисляем

[ dp[u]=(w[u]-x)+\sum_{\text{child }v}\max(0,dp[v]). ]

Положительная ветка полезна, отрицательная нет, поэтому отрицательные значения никогда не присоединяются. 4. Одновременно запоминаем максимальное dp[v] среди детей. Если хотя бы один сын имеет dp[v] >= 0, можно построить допустимое множество, содержащее u и этого сына. При этом все остальные положительные ветки тоже выгодно добавить. Если сумма положительных веток равна нулю, но лучший сын имеет значение ноль, всё равно его нужно добавить, поскольку это увеличивает размер группы с одного до двух, не меняя сумму. 5. Если для некоторой вершины полученная сумма допустимого множества больше нуля, проверка для x возвращает True. Это означает, что существует группа со средним больше x. 6. Выполняем двоичный поиск между минимальной и максимальной силой. Если проверка успешна, двигаем левую границу вверх, иначе двигаем правую границу вниз. После 60 итераций левая граница достаточно близка к настоящему ответу.

Why it works

Рассмотрим любую допустимую группу. После фиксации корня у неё существует единственная вершина u, ближайшая к корню дерева. Все остальные вершины группы лежат в отдельных дочерних поддеревьях u, причём внутри каждого выбранного дочернего поддерева группа обязана быть связной и содержать соответствующего сына. Именно такие варианты описывает dp[v]. Для максимальной суммы отрицательную ветку брать никогда не выгодно, а положительную всегда выгодно брать.

Чтобы группа содержала хотя бы две вершины, для некоторой вершины u должен быть выбран хотя бы один сын. Проверка явно требует такого сына, поэтому одиночная вершина никогда не считается допустимой. Следовательно, проверка возвращает True тогда и только тогда, когда существует допустимая группа с положительной суммой w[v]-x, а это эквивалентно существованию группы со средним больше x. Монотонность этого условия делает двоичный поиск корректным.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    n = int(input())
    w = list(map(int, input().split()))

    g = [[] for _ in range(n)]
    for _ in range(n - 1):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    parent = [-1] * n
    parent[0] = n
    order = [0]

    for u in order:
        for v in g[u]:
            if v == parent[u]:
                continue
            parent[v] = u
            order.append(v)

    dp = [0.0] * n

    def possible(x):
        for u in reversed(order):
            base = w[u] - x
            positive_sum = 0.0
            max_child = -float("inf")

            for v in g[u]:
                if parent[v] != u:
                    continue

                value = dp[v]

                if value > 0.0:
                    positive_sum += value

                if value > max_child:
                    max_child = value

            dp[u] = base + positive_sum

            if max_child >= 0.0:
                candidate = base + positive_sum
                if candidate > 0.0:
                    return True

        return False

    lo = float(min(w))
    hi = float(max(w))

    for _ in range(60):
        mid = (lo + hi) / 2.0
        if possible(mid):
            lo = mid
        else:
            hi = mid

    print(f"{lo:.15f}")

if __name__ == "__main__":
    solve()

Сначала строится обычное неориентированное представление дерева. Затем parent и order позволяют получить направление от корня к листьям без рекурсии. Массив order имеет свойство, что при проходе справа налево дети уже обработаны до своего родителя.

Внутри possible переменная dp[u] хранит лучшую сумму для связного множества, содержащего u, без требования размера хотя бы два. Именно это значение нужно родителю, потому что родитель может захотеть присоединить только одного сына или вообще не присоединять его.

positive_sum содержит сумму всех положительных dp детей. max_child нужен именно для ограничения размера. Если все дети отрицательны, вершина u может дать только одноэлементную группу, и она не должна влиять на ответ. Если существует ребёнок с нулевым dp, его тоже можно присоединить, поэтому проверяется max_child >= 0.0, а не только max_child > 0.0.

Нижняя граница двоичного поиска равна минимальному весу. Любая пара соседних вершин имеет среднее не меньше минимального веса, поэтому настоящий ответ не ниже этой границы. Верхняя граница равна максимальному весу, поскольку среднее не может быть больше максимального элемента.

Используются 60 итераций. Этого достаточно даже для диапазона 10^9, поскольку после 60 делений пополам его длина становится примерно 8.7 * 10^-10. Python использует числа двойной точности, чего достаточно для требуемой абсолютной или относительной погрешности.

Worked Examples

Рассмотрим Sample 1:

Корнем является вершина 1, поэтому дерево направлено как 1 -> 2 -> 3. Истинный ответ равен 2.5. Проверим два значения по разные стороны от него.

`x`	`dp[3]`	`dp[2]`	`dp[1]`	Допустимая группа с положительной суммой
`2.4`	`0.6`	`0.2`	`-1.2`	`{2,3}`, сумма `0.2`, `True`
`2.6`	`0.4`	`-0.2`	`-1.6`	Нет, `False`

При x = 2.4 вершины 2 и 3 после преобразования имеют значения -0.4 и 0.6. Их сумма равна 0.2, то есть их исходное среднее 2.5 действительно больше 2.4. При x = 2.6 даже эта лучшая пара даёт отрицательную сумму, поэтому значение уже недостижимо.

Рассмотрим Sample 2:

Здесь максимальное среднее равно 5. Группа из всех трёх вершин имеет среднее ровно 5, поэтому для двоичного поиска особенно показательно проверить значения чуть ниже и чуть выше ответа.

`x`	`dp[3]`	`dp[2]`	`dp[1]`	Результат
`4.9`	`2.1`	`-1.8`	`0.3`	`True`, группа `{2,3}`
`5.1`	`1.9`	`-2.2`	`1.9`	`False`

При x = 4.9 преобразованные веса равны 2.1, -3.9, 2.1. Пара {2,3} имеет сумму -1.8, но для вершины 2 она всё равно позволяет построить допустимую группу {2,3} с суммой 0.3 после правильного вычисления: -3.9 + 2.1 = -1.8. Здесь положительной группой является не эта пара, а вершина 1 вместе с её поддеревом только если поддерево имеет положительный вклад. В данном конкретном случае dp[1] = 2.1 + (-1.8) = 0.3, но это одноэлементная группа относительно DP. Поэтому для строгого требования размера этот результат не должен использоваться как окончательное свидетельство.

Этот пример показывает тонкость реализации: положительный dp[u] сам по себе недостаточен. Для x = 4.9 допустимой положительной группой является {1,2,3}, чья сумма равна 0.3, потому что родитель 1 присоединяет ребёнка 2, даже несмотря на отрицательное значение dp[2]. Это указывает на проблему в упрощённой проверке, если она рассматривает только положительные dp детей.

Для корректной реализации нам нужно учитывать, что отрицательная dp ветка иногда необходима, чтобы получить группу размера хотя бы два, даже если сама по себе она уменьшает сумму. Именно поэтому окончательная проверка должна отдельно рассматривать лучший вариант с обязательным ребёнком.

Corrected DP Detail

Чтобы избежать указанной ловушки, введём ещё одно значение best2[u]. Оно означает максимальную сумму связного множества, которое содержит u и содержит хотя бы одного ребёнка u.

Для dp[u] всё остаётся прежним:

[ dp[u]=a_u+\sum_v\max(0,dp[v]). ]

Для best2[u] сначала можно взять все положительные ветви. Если среди них есть хотя бы одна, этого уже достаточно для размера хотя бы два. Если положительных ветвей нет, всё равно нужно выбрать лучшего ребёнка, даже если его dp отрицателен:

a_u+ \begin{cases} \sum_v\max(0,dp[v]), & \text{если существует }dp[v]>0,
\max_v dp[v], & \text{иначе}. \end{cases} ]

Именно best2[u] > 0 является правильным условием успеха.

Следовательно, приведённый выше код должен использовать max_child без условия max_child >= 0. Нужно проверять и отрицательного ребёнка, если положительных детей нет. Ниже находится окончательная версия решения.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    n = int(input())
    w = list(map(int, input().split()))

    g = [[] for _ in range(n)]
    for _ in range(n - 1):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    parent = [-1] * n
    parent[0] = n
    order = [0]

    for u in order:
        for v in g[u]:
            if v == parent[u]:
                continue
            parent[v] = u
            order.append(v)

    dp = [0.0] * n

    def possible(x):
        for u in reversed(order):
            base = w[u] - x
            positive_sum = 0.0
            max_child = -float("inf")

            for v in g[u]:
                if parent[v] != u:
                    continue

                value = dp[v]

                if value > 0.0:
                    positive_sum += value

                if value > max_child:
                    max_child = value

            dp[u] = base + positive_sum

            if max_child != -float("inf"):
                if positive_sum > 0.0:
                    best_with_child = base + positive_sum
                else:
                    best_with_child = base + max_child

                if best_with_child > 0.0:
                    return True

        return False

    lo = float(min(w))
    hi = float(max(w))

    for _ in range(60):
        mid = (lo + hi) / 2.0
        if possible(mid):
            lo = mid
        else:
            hi = mid

    print(f"{lo:.15f}")

if __name__ == "__main__":
    solve()

Ключевое исправление находится в вычислении best_with_child. Если существует положительная дочерняя ветка, мы берём все положительные ветки. Если положительных веток нет, мы обязаны выбрать хотя бы одну, и тогда оптимальным выбором является ребёнок с максимальным dp, даже если это значение отрицательное.

Это особенно важно на цепочке 7, 1, 7. При проверке значения чуть меньше 5 группа из всех трёх вершин положительна после преобразования, хотя промежуточная динамика для вершины 2 остаётся отрицательной. Если разрешать только положительные дочерние dp, такая группа была бы потеряна.

Complexity Analysis

Measure	Complexity	Explanation
Time	`O(n log(10^9 / eps))`	Каждая из 60 итераций двоичного поиска делает один проход по всем вершинам и рёбрам
Space	`O(n)`	Дерево, родители, порядок обхода и массив DP занимают линейную память

При n = 200000 одна проверка имеет линейную сложность, а всего выполняется 60 проверок. Это около нескольких десятков миллионов простых операций по рёбрам, что существенно меньше квадратичного перебора порядка 10^10 вариантов. Итеративный обход также позволяет безопасно работать с деревом, являющимся длинной цепочкой.

Test Cases

import sys
import io

def solve():
    input = sys.stdin.readline

    n = int(input())
    w = list(map(int, input().split()))

    g = [[] for _ in range(n)]
    for _ in range(n - 1):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        g[a].append(b)
        g[b].append(a)

    parent = [-1] * n
    parent[0] = n
    order = [0]

    for u in order:
        for v in g[u]:
            if v == parent[u]:
                continue
            parent[v] = u
            order.append(v)

    dp = [0.0] * n

    def possible(x):
        for u in reversed(order):
            base = w[u] - x
            positive_sum = 0.0
            max_child = -float("inf")

            for v in g[u]:
                if parent[v] != u:
                    continue

                value = dp[v]

                if value > 0.0:
                    positive_sum += value

                if value > max_child:
                    max_child = value

            dp[u] = base + positive_sum

            if max_child != -float("inf"):
                if positive_sum > 0.0:
                    best_with_child = base + positive_sum
                else:
                    best_with_child = base + max_child

                if best_with_child > 0.0:
                    return True

        return False

    lo = float(min(w))
    hi = float(max(w))

    for _ in range(60):
        mid = (lo + hi) / 2.0
        if possible(mid):
            lo = mid
        else:
            hi = mid

    print(f"{lo:.15f}")

def run(inp: str) -> str:
    global_input = sys.stdin
    global_output = sys.stdout

    try:
        sys.stdin = io.StringIO(inp)
        sys.stdout = io.StringIO()
        solve()
        return sys.stdout.getvalue().strip()
    finally:
        sys.stdin = global_input
        sys.stdout = global_output

assert run("""3
1 2 3
1 2
2 3
""") == "2.500000000000000", "sample 1"

assert run("""3
7 1 7
1 2
2 3
""") == "5.000000000000000", "sample 2"

assert run("""4
7 1 7 7
1 2
2 3
2 4
""") == "5.500000000000000", "sample 3"

assert run("""2
0 1
1 2
""") == "0.500000000000000", "minimum size"

assert run("""5
42 42 42 42 42
1 2
2 3
3 4
4 5
""") == "42.000000000000000", "all equal"

assert run("""4
10 0 0 10
1 2
2 3
3 4
""") == "5.000000000000000", "disconnected maximum values"

n = 200000
weights = " ".join(["1000000000"] * n)
edges = "\n".join(f"{i} {i + 1}" for i in range(1, n))
max_case = f"{n}\n{weights}\n{edges}\n"

assert run(max_case) == "1000000000.000000000000000", "maximum size"

Test input	Expected output	What it validates
`2 / 0 1 / 1 2`	`0.500000000000000`	Минимальный размер, группа обязана содержать обе вершины
`5 / 42 42 42 42 42`	`42.000000000000000`	Все веса одинаковы, ответ совпадает с весом любой вершины
`4 / 10 0 0 10`	`5.000000000000000`	Нельзя выбрать две несвязанные сильные вершины
Цепочка из `200000` вершин с весами `10^9`	`1000000000.000000000000000`	Максимальный размер и верхняя граница весов

Edge Cases

Случай с двумя вершинами проверяет сразу условие минимального размера группы. Для входа

2
100 0
1 2

при x = 75 преобразованные веса равны 25 и -75. Для первой вершины её собственное dp положительно, но это одиночная вершина, поэтому она не может дать True. Чтобы получить допустимую группу, нужно добавить вторую вершину, после чего сумма становится -50. Проверка возвращает False, а двоичный поиск сходится к 50.

Случай с несвязанными сильными вершинами показывает, почему DP должно учитывать структуру дерева. Для

вершины 1 и 4 обе имеют вес 10, но между ними находятся две вершины нулевого веса. Группа {1,4} недопустима. Максимум достигается на группе {1,2} или {3,4} и равен 5. DP автоматически учитывает промежуточные вершины, потому что любое выбранное множество строится из связных дочерних поддеревьев.

Случай, где оптимальная группа больше двух вершин, особенно хорошо проверяет корректность динамики. Для

каждое отдельное ребро имеет среднее 5, но вся группа имеет сумму 30 и размер 4, то есть среднее 7.5. При проверке x < 7.5 все три дочерние ветви дают положительный вклад, и DP вершины 1 присоединяет их одновременно. Это показывает, почему нельзя ограничиваться выбором одной лучшей дочерней ветки.

Наконец, значение dp ребёнка может быть отрицательным, хотя этот ребёнок всё равно обязан быть выбран для получения группы размера хотя бы два. Именно это происходит в цепочке

при x = 4.9. Для вершины 2 значение равно -1.8, потому что 1 - 4.9 + (7 - 4.9) = -1.8. Но вершина 1 имеет преобразованный вес 2.1, поэтому группа {1,2,3} имеет суммарный преобразованный вес 0.3 и действительно улучшает ответ относительно 4.9. Отдельное хранение dp и лучшего варианта с обязательным ребёнком позволяет не потерять такие группы.

Важная деталь: я исправил тонкость в DP после первого варианта проверки. Для требования «хотя бы две вершины» отрицательный dp ребёнка иногда всё равно необходимо присоединить, поэтому проверка только неотрицательных дочерних значений была бы некорректной.