結果

問題 No.1917 LCMST
ユーザー qwewe
提出日時 2025-05-14 13:21:01
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 3,421 ms / 4,000 ms
コード長 9,691 bytes
コンパイル時間 441 ms
コンパイル使用メモリ 82,432 KB
実行使用メモリ 301,384 KB
最終ジャッジ日時 2025-05-14 13:23:33
合計ジャッジ時間 58,032 ms
ジャッジサーバーID
(参考情報)
judge4 / judge5
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff #

import math
import sys

# Use sys.stdin.readline for faster input reading
# This is important for performance in competitive programming with large inputs
input = sys.stdin.readline

# Function to compute Greatest Common Divisor (GCD)
# We use math.gcd as it's implemented in C and generally faster than a Python implementation.
def gcd(a, b):
    """Computes the Greatest Common Divisor (GCD) of a and b."""
    return math.gcd(a, b)

# Function to compute Least Common Multiple (LCM)
def lcm(a, b):
    """Computes the Least Common Multiple (LCM) of a and b."""
    # Handle cases where a or b is 0. The LCM is typically defined as 0 in this case.
    if a == 0 or b == 0:
        return 0
    
    # Computes LCM using the formula: lcm(a, b) = |a * b| / gcd(a, b)
    # Python's arbitrary precision integers handle potentially large intermediate products |a * b|.
    # Using // ensures integer division.
    # abs() ensures positivity, although inputs A_i are positive.
    val = abs(a * b) // gcd(a, b)
    return val

def solve():
    """Solves the LCMST problem."""
    
    # Read number of vertices N
    N = int(input())
    # Read the list of values A_i associated with vertices
    A = list(map(int, input().split()))

    has_one = False
    total_sum = 0
    # Dictionary to store counts of each distinct value A_i
    A_counts = {}
    
    # Calculate total sum of A_i values and check if the value 1 is present
    # Also populate the A_counts dictionary to count occurrences of each value
    for x in A:
        if x == 1:
            has_one = True
        total_sum += x
        A_counts[x] = A_counts.get(x, 0) + 1
    
    # Special Case: If any A_i is 1
    # If a vertex k exists with A_k=1, the Minimum Spanning Tree (MST) can be formed
    # by connecting vertex k to all other vertices j. The weight of edge (k, j) is
    # lcm(A_k, A_j) = lcm(1, A_j) = A_j.
    # The total weight of this star graph (which is a spanning tree) is sum(A_j for j != k).
    # This sum is equal to (sum(A_i for all i)) - A_k = total_sum - 1 (since A_k=1).
    # This star graph is indeed the MST because any edge (i, j) not involving k has weight 
    # lcm(A_i, A_j) >= max(A_i, A_j), while the path i-k-j in the star graph has maximum edge 
    # weight max(A_i, A_j). By the cycle property of MSTs, the star graph is minimal.
    if has_one:
        print(total_sum - 1)
        return

    # General Case: No A_i is 1
    
    # Get the unique values present in A and sort them. These represent the 'types' of vertices.
    unique_A = sorted(A_counts.keys())
    num_distinct_vals = len(unique_A) # k = number of distinct values

    # Set the maximum possible value for A_i + 1 (used for array/list bounds)
    # The constraint states A_i <= 10^5.
    V_max_limit = 100001 

    # Use a set for efficient checking (`O(1)` average time) if a value exists among the unique A_i values.
    present_set = set(unique_A)
    
    # Map each unique value to a unique index from 0 to k-1. This is needed for the DSU structure.
    val_map = {val: i for i, val in enumerate(unique_A)}

    # Calculate the initial component of the MST weight.
    # If multiple vertices have the same value x, they must be connected in the MST.
    # The cheapest way is to form a path/tree using edges of weight lcm(x, x) = x.
    # To connect `count` vertices, we need `count - 1` edges. Total cost for value x is x * (count - 1).
    initial_mst_weight = 0
    for x in unique_A:
        count = A_counts[x]
        if count > 1:
            initial_mst_weight += x * (count - 1)

    # List to store candidate edges for Kruskal's algorithm. Each element is a tuple (weight, u, v).
    candidate_edges = []

    # Precompute lists of multiples for the GCD iteration method.
    # `multiples_data[g]` will store a list of integers `x` such that `g*x` is one of the unique values in A.
    # This precomputation helps efficiently find pairs for edge generation.
    multiples_data = [[] for _ in range(V_max_limit)] 
    # Iterate through all possible GCD values `g` from 1 up to V_max_limit.
    for g in range(1, V_max_limit):
        # Iterate through multiples of `g` starting from `g` itself.
        for multiple in range(g, V_max_limit, g):
            # Check if this multiple `g*x` is present in the input values `A`.
            if multiple in present_set:
                 # If yes, store `x = multiple // g` in the list associated with `g`.
                 multiples_data[g].append(multiple // g) 

    # Generate candidate edges using the GCD iteration logic.
    # This technique is derived from solutions to similar problems and is efficient.
    # It generates a limited set of edges sufficient to find the MST.
    for g in range(1, V_max_limit):
        # Get the list of factors `x` for the current GCD `g`.
        # The list `S_g` is guaranteed to be sorted because the inner loop iterates `multiple` in increasing order.
        S_g = multiples_data[g] 
        # If there are at least two multiples of `g` present in A (i.e., |S_g| >= 2)
        if len(S_g) >= 2:
            # Find the minimum factor `x_min` among these multiples.
            x_min = S_g[0]
            val_min = g * x_min # The actual vertex value corresponding to x_min
            
            # Consider edges connecting `val_min` to all other values `g*x` where `x` is in `S_g`.
            for i in range(1, len(S_g)):
                x = S_g[i]
                val_curr = g * x # The actual vertex value corresponding to x
                
                # Calculate the LCM, which is the weight of the edge (val_min, val_curr).
                w = lcm(val_min, val_curr)
                # Add this potential edge (weight, endpoint1, endpoint2) to the list of candidates.
                candidate_edges.append((w, val_min, val_curr))

    # Sort all generated candidate edges by weight in non-decreasing order.
    # This is a crucial step for Kruskal's algorithm.
    candidate_edges.sort()

    # Initialize Disjoint Set Union (DSU) structure.
    # Each unique value `A_i` initially represents a separate component/set.
    # `parent[i]` stores the parent of element `i`. Initially, each element is its own parent.
    parent = list(range(num_distinct_vals))
    # `size[i]` stores the size of the tree rooted at `i`. Used for Union by Size optimization.
    size = [1] * num_distinct_vals 

    # DSU Find operation with Path Compression optimization.
    # Finds the representative (root) of the set containing element `i`.
    def find(i):
        # Find the root by traversing up the parent pointers.
        root = i
        while parent[root] != root:
             root = parent[root]
        # Path compression: Make all nodes on the path from `i` to the root point directly to the root.
        # This flattens the tree structure, speeding up future Find operations.
        curr = i
        while parent[curr] != root:
             next_node = parent[curr]
             parent[curr] = root
             curr = next_node
        return root

    # DSU Unite operation with Union by Size optimization.
    # Merges the sets containing elements `i` and `j`.
    def unite(i, j):
        # Find the roots of the sets containing `i` and `j`.
        root_i = find(i)
        root_j = find(j)
        # If `i` and `j` are already in the same set, do nothing.
        if root_i != root_j:
            # Union by Size: Attach the smaller tree to the root of the larger tree.
            # This helps keep the tree depth low, improving efficiency.
            if size[root_i] < size[root_j]:
                root_i, root_j = root_j, root_i # Ensure root_i is the root of the larger tree
            # Make the root of the smaller tree (`root_j`) point to the root of the larger tree (`root_i`).
            parent[root_j] = root_i
            # Update the size of the merged tree rooted at `root_i`.
            size[root_i] += size[root_j]
            return True # Return True indicating that a merge occurred.
        return False # Return False if `i` and `j` were already in the same set.

    # Kruskal's algorithm main loop.
    # Starts with the initial weight from connecting identical value nodes.
    current_mst_weight = initial_mst_weight
    # Counter for the number of edges added to the MST.
    edges_count = 0 
    
    # If there's only one distinct value type (e.g., all A_i are the same), the MST cost
    # is just the initial weight connecting these identical nodes.
    if num_distinct_vals <= 1:
         print(initial_mst_weight) 
         return

    # Process the candidate edges, sorted by weight.
    for w, u, v in candidate_edges:
        # Map the vertex values `u` and `v` to their corresponding DSU indices (0 to k-1).
        u_idx = val_map[u]
        v_idx = val_map[v]
        
        # Try to unite the sets containing `u` and `v`.
        # The `unite` function returns `True` if `u` and `v` were in different components
        # and have now been merged. It returns `False` if they were already connected.
        if unite(u_idx, v_idx): 
            # If merged, this edge is part of the MST. Add its weight to the total cost.
            current_mst_weight += w 
            # Increment the count of edges added to the MST.
            edges_count += 1 
            # The MST for `k` components requires exactly `k-1` edges.
            # If we have added `k-1` edges, the MST is complete, so we can stop.
            if edges_count == num_distinct_vals - 1:
                break 

    # Print the final calculated total weight of the Minimum Spanning Tree.
    print(current_mst_weight)

# Execute the main solve function to run the program.
solve()
0