Parametric MST (Codeforces 1656F)
Tags: Greedy, graphs

You are given $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$. For any real number $$$t$$$, consider the complete weighted graph on $$$n$$$ vertices $$$K_n(t)$$$ with weight of the edge between vertices $$$i$$$ and $$$j$$$ equal to $$$w_{ij}(t) = a_i \cdot a_j + t \cdot (a_i + a_j)$$$. Let $$$f(t)$$$ be the cost of the minimum spanning tree of $$$K_n(t)$$$. Determine whether $$$f(t)$$$ is bounded above and, if so, output the maximum value it attains.

Input: The input consists of multiple test cases. The first line contains a single integer $$$T$$$ ($$$1 \leq T \leq 10^4$$$) — the number of test cases. Description of the test cases follows.The first line of each test case contains an integer $$$n$$$ ($$$2 \leq n \leq 2 \cdot 10^5$$$) — the number of vertices of the graph.The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$-10^6 \leq a_i \leq 10^6$$$).The sum of $$$n$$$ for all test cases is at most $$$2 \cdot 10^5$$$.

Output: For each test case, print a single line with the maximum value of $$$f(t)$$$ (it can be shown that it is an integer), or INF if $$$f(t)$$$ is not bounded above. 


Sample Inputs:
5
2
1 0
2
-1 1
3
1 -1 -2
3
3 -1 -2
4
1 2 3 -4

Sample Outputs:
INF
-1
INF
-6
-18

Solution (Python):
import io
import os
FILE_INPUT_MODE = 0
if FILE_INPUT_MODE >= 0:
    if FILE_INPUT_MODE == 1:
        CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
        FD = os.open(f'{CURRENT_DIR}\\input.txt', os.O_RDONLY)
    elif FILE_INPUT_MODE == 0:
        FD = 0
    DATA = io.BytesIO(os.read(FD, os.fstat(FD).st_size))
    if FD != 0:
        os.close(FD)

    def input():
        return DATA.readline().decode().rstrip('\r\n')

def get(f, t):
    return f[0] + f[1] * t

def solve(n, a):
    a.sort()
    cur = [0, 0]
    for i in range(n - 1):
        cur[0] += a[i] * a[n - 1]
        cur[1] += a[i] + a[n - 1]
    if cur[1] < 0:
        return 'INF'
    ans = min(get(cur, -a[0]), get(cur, -a[n - 1]))
    for i in range(n - 2, 0, -1):
        t = -a[i]
        ans = max(ans, get(cur, t))
        cur[0] -= a[i] * a[n - 1]
        cur[1] -= a[i] + a[n - 1]
        cur[0] += a[i] * a[0]
        cur[1] += a[i] + a[0]
    if cur[1] > 0:
        return 'INF'
    return ans

def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = [int(word) for word in input().strip().split()]
        print(solve(n, a))
if __name__ == '__main__':
    main()

Solution (Python):
import io
import os
FILE_INPUT_MODE = 0
if FILE_INPUT_MODE >= 0:
    if FILE_INPUT_MODE == 1:
        CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
        FD = os.open(f'{CURRENT_DIR}\\input.txt', os.O_RDONLY)
    elif FILE_INPUT_MODE == 0:
        FD = 0
    DATA = io.BytesIO(os.read(FD, os.fstat(FD).st_size))
    if FD != 0:
        os.close(FD)

    def input():
        return DATA.readline().decode().rstrip('\r\n')

def get(f, t):
    return f[0] + f[1] * t

def solve(n, a):
    a.sort()
    ans = -float('inf')
    cur = [0, 0]
    for i in range(n - 1):
        cur[0] += a[i] * a[n - 1]
        cur[1] += a[i] + a[n - 1]
    if cur[1] < 0:
        return 'INF'
    for i in range(n - 1, -1, -1):
        t = -a[i]
        ans = max(ans, get(cur, t))
        if i != n - 1 and i != 0:
            cur[0] -= a[i] * a[n - 1]
            cur[1] -= a[i] + a[n - 1]
            cur[0] += a[i] * a[0]
            cur[1] += a[i] + a[0]
    if cur[1] > 0:
        return 'INF'
    return ans

def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = [int(word) for word in input().strip().split()]
        print(solve(n, a))
if __name__ == '__main__':
    main()

Solution (Python):
import os
import sys
from io import BytesIO, IOBase
_str = str
str = lambda x=b'': x if type(x) is bytes else _str(x).encode()
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
from math import inf

def solve(a, n):
    a = sorted(a)
    mul_sum = 0
    add_sum = 0
    for i in range(1, n):
        mul_sum += a[0] * a[i]
        add_sum += a[0] + a[i]
    res = mul_sum + add_sum * -a[0]
    if add_sum > 0:
        return 'INF'
    for i in range(1, n - 1):
        mul_sum += a[-1] * a[i] - a[0] * a[i]
        add_sum += a[-1] - a[0]
        res = max(res, mul_sum + add_sum * -a[i])
    if add_sum < 0:
        return 'INF'
    return res

def main():
    t = int(input())
    while t:
        n = int(input())
        a = list(map(int, input().strip().split(' ')))
        print(solve(a, n))
        t -= 1
    return
main()