logoChat01.ai

질문

image.png
image.png
image.png
image.png

Give me the python3 code

Python

python
def findPrice(tree_nodes, tree_from, tree_to, values):
    import sys
    sys.setrecursionlimit(1_000_000)

    adj = [[] for _ in range(tree_nodes + 1)]
    for u, v in zip(tree_from, tree_to):
        adj[u].append(v)
        adj[v].append(u)

    ans = [0] * tree_nodes

    def dfs(u, parent):
        candidates = {u: values[u - 1]}
        is_leaf = True

        for v in adj[u]:
            if v == parent:
                continue
            is_leaf = False
            child_max, child_min = dfs(v, u)

            for val, node in child_max:
                candidates[node] = val
            for val, node in child_min:
                candidates[node] = val

        max_pairs = []  # sorted by value descending
        min_pairs = []  # sorted by value ascending

        for node, val in candidates.items():
            pair = (val, node)

            if len(max_pairs) < 2:
                max_pairs.append(pair)
                max_pairs.sort(key=lambda x: x[0], reverse=True)
            elif val > max_pairs[-1][0]:
                max_pairs[-1] = pair
                if max_pairs[0][0] < max_pairs[1][0]:
                    max_pairs[0], max_pairs[1] = max_pairs[1], max_pairs[0]

            if len(min_pairs) < 2:
                min_pairs.append(pair)
                min_pairs.sort(key=lambda x: x[0])
            elif val < min_pairs[-1][0]:
                min_pairs[-1] = pair
                if min_pairs[0][0] > min_pairs[1][0]:
                    min_pairs[0], min_pairs[1] = min_pairs[1], min_pairs[0]

        if is_leaf:
            ans[u - 1] = 0
        else:
            best = max_pairs[0][0] * max_pairs[1][0]
            best = max(best, min_pairs[0][0] * min_pairs[1][0])
            ans[u - 1] = best

        return max_pairs, min_pairs

    dfs(1, 0)
    return ans

print(findPrice(6, [1,2,2,5], [2,3,4,6], [3,4,-2,0,-1,6]))
print(findPrice(6, [1,2,2,2,4], [2,3,4,5,6], [1,0,-1,-2,-3,9]))
print(findPrice(4, [1,2,3], [2,3,4], [-7,-4,-2,3]))

Python

python
print(findPrice(6, [1,2,2,1,5], [2,3,4,5,6], [3,4,-2,0,-1,6]))

사고

Clarifying value mismatch and algorithm check

The sample output seemed off at first, but after re-checking the values and understanding the input format, it makes sense: expected node values are [-7, -4, 2, -3], not [-7, -4, -2, 3]. I now see how the algorithm reaches the correct sample output. I'll now finalize the code and add a brief explanation, keeping the complexity at O(n).

Thought for 2m 25s

답변

python
#!/bin/python3

import os
import sys

sys.setrecursionlimit(1_000_000)


def findPrice(tree_nodes, tree_from, tree_to, values):
    # Build the tree as an undirected graph, then root it at 1
    graph = [[] for _ in range(tree_nodes + 1)]
    for u, v in zip(tree_from, tree_to):
        graph[u].append(v)
        graph[v].append(u)

    ans = [0] * tree_nodes

    def add_max(best, pair):
        # keep top 2 largest (value, node)
        if len(best) < 2:
            best.append(pair)
            if len(best) == 2 and best[0][0] < best[1][0]:
                best[0], best[1] = best[1], best[0]
        elif pair[0] > best[1][0]:
            best[1] = pair
            if best[0][0] < best[1][0]:
                best[0], best[1] = best[1], best[0]

    def add_min(best, pair):
        # keep top 2 smallest (value, node)
        if len(best) < 2:
            best.append(pair)
            if len(best) == 2 and best[0][0] > best[1][0]:
                best[0], best[1] = best[1], best[0]
        elif pair[0] < best[1][0]:
            best[1] = pair
            if best[0][0] > best[1][0]:
                best[0], best[1] = best[1], best[0]

    def dfs(u, parent):
        # Candidate extreme nodes from this subtree.
        # key = node id, value = node value
        candidates = {u: values[u - 1]}
        is_leaf = True

        for v in graph[u]:
            if v == parent:
                continue
            is_leaf = False
            child_max2, child_min2 = dfs(v, u)

            # Only child's 2 largest and 2 smallest are enough
            for val, node_id in child_max2:
                candidates[node_id] = val
            for val, node_id in child_min2:
                candidates[node_id] = val

        max2 = []
        min2 = []

        for node_id, val in candidates.items():
            pair = (val, node_id)
            add_max(max2, pair)
            add_min(min2, pair)

        if is_leaf:
            ans[u - 1] = 0
        else:
            prod1 = max2[0][0] * max2[1][0]
            prod2 = min2[0][0] * min2[1][0]
            ans[u - 1] = max(prod1, prod2)

        return max2, min2

    dfs(1, 0)
    return ans


if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    tree_nodes, tree_edges = map(int, input().split())

    tree_from = []
    tree_to = []

    for _ in range(tree_edges):
        u, v = map(int, input().split())
        tree_from.append(u)
        tree_to.append(v)

    values_count = int(input().strip())
    values = [int(input().strip()) for _ in range(values_count)]

    result = findPrice(tree_nodes, tree_from, tree_to, values)

    fptr.write('\n'.join(map(str, result)))
    fptr.write('\n')
    fptr.close()

This works in O(n) time and handles the negative-value cases correctly.

이 Q&A 공유하기