Spanning Tree Queries (Codeforces 1633E)
Tags: Greedy, dsu

You are given a connected weighted undirected graph, consisting of $$$n$$$ vertices and $$$m$$$ edges.You are asked $$$k$$$ queries about it. Each query consists of a single integer $$$x$$$. For each query, you select a spanning tree in the graph. Let the weights of its edges be $$$w_1, w_2, \dots, w_{n-1}$$$. The cost of a spanning tree is $$$\sum \limits_{i=1}^{n-1} |w_i - x|$$$ (the sum of absolute differences between the weights and $$$x$$$). The answer to a query is the lowest cost of a spanning tree.The queries are given in a compressed format. The first $$$p$$$ $$$(1 \le p \le k)$$$ queries $$$q_1, q_2, \dots, q_p$$$ are provided explicitly. For queries from $$$p+1$$$ to $$$k$$$, $$$q_j = (q_{j-1} \cdot a + b) \mod c$$$.Print the xor of answers to all queries.

Input: The first line contains two integers $$$n$$$ and $$$m$$$ ($$$2 \le n \le 50$$$; $$$n - 1 \le m \le 300$$$) — the number of vertices and the number of edges in the graph.Each of the next $$$m$$$ lines contains a description of an undirected edge: three integers $$$v$$$, $$$u$$$ and $$$w$$$ ($$$1 \le v, u \le n$$$; $$$v \neq u$$$; $$$0 \le w \le 10^8$$$) — the vertices the edge connects and its weight. Note that there might be multiple edges between a pair of vertices. The edges form a connected graph.The next line contains five integers $$$p, k, a, b$$$ and $$$c$$$ ($$$1 \le p \le 10^5$$$; $$$p \le k \le 10^7$$$; $$$0 \le a, b \le 10^8$$$; $$$1 \le c \le 10^8$$$) — the number of queries provided explicitly, the total number of queries and parameters to generate the queries.The next line contains $$$p$$$ integers $$$q_1, q_2, \dots, q_p$$$ ($$$0 \le q_j < c$$$) — the first $$$p$$$ queries.

Output: Print a single integer — the xor of answers to all queries.

Note: The queries in the first example are $$$0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0$$$. The answers are $$$11, 9, 7, 3, 1, 5, 8, 7, 5, 7, 11$$$.  The queries in the second example are $$$3, 0, 2, 1, 6, 0, 3, 5, 4, 1$$$. The answers are $$$14, 19, 15, 16, 11, 19, 14, 12, 13, 16$$$.  The queries in the third example are $$$75, 0, 0, \dots$$$. The answers are $$$50, 150, 150, \dots$$$.  

Sample Inputs:
5 8
4 1 4
3 1 0
3 5 3
2 5 4
3 4 8
4 3 4
4 2 8
5 3 9
3 11 1 1 10
0 1 2

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

3 3
1 2 50
2 3 100
1 3 150
1 10000000 0 0 100000000
75

Sample Outputs:
4

5

164

Solution (Python):
class UnionFind:

    def __init__(self, n):
        self.parent = list(range(n))

    def find(self, a):
        acopy = a
        while a != self.parent[a]:
            a = self.parent[a]
        while acopy != a:
            (self.parent[acopy], acopy) = (a, self.parent[acopy])
        return a

    def union(self, a, b):
        self.parent[self.find(b)] = self.find(a)

def calculateSignature(edges, x, n):
    """"""
    (lo, hi) = (-1, len(edges))
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if edges[mid][2] <= x:
            lo = mid
        else:
            hi = mid - 1
    uf = UnionFind(n + 1)
    cost = 0
    takenEdgesSerial = 1
    leftCnts = rightCnts = 0
    (l, r) = (lo, lo + 1)
    cnts = 0
    while cnts < n - 1:
        if abs(x - edges[l][2]) < abs(x - edges[r][2]):
            idx = l
            l -= 1
        else:
            idx = r
            r += 1
        (u, v, w) = edges[idx]
        if uf.find(u) != uf.find(v):
            uf.union(u, v)
            cost += abs(x - w)
            if w < x:
                leftCnts += 1
            if w > x:
                rightCnts += 1
            takenEdgesSerial *= 1 + idx
            takenEdgesSerial %= 1000000007
            cnts += 1
    return (cost, takenEdgesSerial, leftCnts, rightCnts)

def main():
    (n, m) = readIntArr()
    edges = []
    for i in range(m):
        (u, v, w) = readIntArr()
        edges.append((u, v, w))
    minXes = []
    costs = []
    leftCnts = []
    rightCnts = []
    edges.append((-1, -1, -int(2000000000.0)))
    edges.append((-1, -1, int(2000000000.0)))
    edges.sort(key=lambda data: data[2])
    R = int(100000000.0 + 5)
    l = 0
    while l < R:
        (cost, tE, lc, rc) = calculateSignature(edges, l, n)
        minXes.append(l)
        costs.append(cost)
        leftCnts.append(lc)
        rightCnts.append(rc)
        (lo, hi) = (l - 1, R)
        while lo < hi:
            mid = (lo + hi + 1) // 2
            (c2, te2, lc2, rc2) = calculateSignature(edges, mid, n)
            if tE == te2 and lc == lc2 and (rc == rc2):
                lo = mid
            else:
                hi = mid - 1
        l = lo + 1
    (p, k, a, b, c) = readIntArr()
    queries = readIntArr()
    for _ in range(k - p):
        queries.append((queries[-1] * a + b) % c)
    xor = 0
    n2 = len(minXes)
    for x in queries:
        (lo, hi) = (-1, n2 - 1)
        while lo < hi:
            mid = (lo + hi + 1) // 2
            if minXes[mid] <= x:
                lo = mid
            else:
                hi = mid - 1
        assert minXes[lo] <= x
        cost = costs[lo] + (x - minXes[lo]) * (leftCnts[lo] - rightCnts[lo])
        xor ^= cost
    print(xor)
    return
import sys
input = sys.stdin.buffer.readline

def oneLineArrayPrint(arr):
    print(' '.join([str(x) for x in arr]))

def multiLineArrayPrint(arr):
    print('\n'.join([str(x) for x in arr]))

def multiLineArrayOfArraysPrint(arr):
    print('\n'.join([' '.join([str(x) for x in y]) for y in arr]))

def readIntArr():
    return [int(x) for x in input().split()]

def makeArr(defaultValFactory, dimensionArr):
    dv = defaultValFactory
    da = dimensionArr
    if len(da) == 1:
        return [dv() for _ in range(da[0])]
    else:
        return [makeArr(dv, da[1:]) for _ in range(da[0])]

def queryInteractive(a, b, c):
    print('? {} {} {}'.format(a, b, c))
    sys.stdout.flush()
    return int(input())

def answerInteractive(x1, x2):
    print('! {} {}'.format(x1, x2))
    sys.stdout.flush()
inf = float('inf')
from math import gcd, floor, ceil
import math
for _abc in range(1):
    main()

Solution (Python):
import io
import os
from bisect import bisect_right
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')

class DisjointSet:

    def __init__(self, n):
        self.fa = [i for i in range(n + 1)]
        self.size = [1] * (n + 1)

    def add(self, u, v):
        u = self.father(u)
        v = self.father(v)
        if u == v:
            return False
        if self.size[u] <= self.size[v]:
            self.fa[u] = v
            self.size[v] += self.size[u]
        else:
            self.fa[v] = u
            self.size[u] += self.size[v]
        return True

    def query(self, u, v):
        return self.father(u) == self.father(v)

    def father(self, u):
        if self.fa[u] == u:
            return u
        v = self.fa[u]
        self.size[v] -= self.size[u]
        w = self.father(v)
        self.size[w] += self.size[u]
        self.fa[u] = w
        return w

def kruskal(n, m, edge, x=0):
    djs = DisjointSet(n)
    i = 0
    while i < m and edge[i][-1] <= x:
        i += 1
    j = i - 1
    mst = 0
    while i < m or j >= 0:
        if j >= 0 and (i >= m or abs(edge[j][-1] - x) <= abs(edge[i][-1] - x)):
            if djs.add(edge[j][0], edge[j][1]):
                mst += abs(edge[j][-1] - x)
            j -= 1
        else:
            if djs.add(edge[i][0], edge[i][1]):
                mst += abs(edge[i][-1] - x)
            i += 1
    return mst

def get_answer(node, mst, q):
    right = bisect_right(node, q)
    left = right - 1
    return ((node[right] - q) * mst[left] + (q - node[left]) * mst[right]) // (node[right] - node[left])

def solve():
    (n, m) = [int(word) for word in input().strip().split()]
    edge = []
    for i in range(m):
        (u, v, w) = [int(word) for word in input().strip().split()]
        edge.append((u, v, w))
    (p, k, a, b, c) = [int(word) for word in input().strip().split()]
    q = [int(word) for word in input().strip().split()]
    edge.sort(key=lambda e: e[-1])
    node = set([0, c])
    for i in range(m):
        node.add(edge[i][-1])
        for j in range(i + 1, m):
            w = edge[i][-1] + edge[j][-1]
            node.add((w + 1) // 2)
            node.add((w + 1) // 2 - 1)
    node = [x for x in node if x >= 0 and x <= c]
    arr = sorted(node)
    node = []
    mst = []
    for x in arr:
        res = kruskal(n, m, edge, x)
        if len(node) > 1:
            delta = mst[-1] - mst[-2]
            if delta % (node[-1] - node[-2]) == 0:
                delta = (mst[-1] - mst[-2]) // (node[-1] - node[-2])
                if mst[-1] + delta * (x - node[-1]) == res:
                    node.pop()
                    mst.pop()
        node.append(x)
        mst.append(res)
    ans = 0
    for i in range(p):
        ans ^= get_answer(node, mst, q[i])
    cur_q = q[-1]
    for i in range(p, k):
        cur_q = (cur_q * a + b) % c
        ans ^= get_answer(node, mst, cur_q)
    return ans

def main():
    print(solve())
if __name__ == '__main__':
    main()

Solution (Python):
import io
import os
from bisect import bisect_right
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')

class DisjointSet:

    def __init__(self, n):
        self.fa = [i for i in range(n + 1)]
        self.size = [1] * (n + 1)

    def add(self, u, v):
        u = self.father(u)
        v = self.father(v)
        if u == v:
            return False
        if self.size[u] <= self.size[v]:
            self.fa[u] = v
            self.size[v] += self.size[u]
        else:
            self.fa[v] = u
            self.size[u] += self.size[v]
        return True

    def query(self, u, v):
        return self.father(u) == self.father(v)

    def father(self, u):
        if self.fa[u] == u:
            return u
        v = self.fa[u]
        w = self.father(v)
        return w

def kruskal(n, m, edge, x=0):
    djs = DisjointSet(n)
    i = 0
    while i < m and edge[i][-1] <= x:
        i += 1
    j = i - 1
    mst = 0
    while i < m or j >= 0:
        if j >= 0 and (i >= m or abs(edge[j][-1] - x) <= abs(edge[i][-1] - x)):
            if djs.add(edge[j][0], edge[j][1]):
                mst += abs(edge[j][-1] - x)
            j -= 1
        else:
            if djs.add(edge[i][0], edge[i][1]):
                mst += abs(edge[i][-1] - x)
            i += 1
    return mst

def get_answer(node, mst, q):
    right = bisect_right(node, q)
    left = right - 1
    return ((node[right] - q) * mst[left] + (q - node[left]) * mst[right]) // (node[right] - node[left])

def solve():
    (n, m) = [int(word) for word in input().strip().split()]
    edge = []
    for i in range(m):
        (u, v, w) = [int(word) for word in input().strip().split()]
        edge.append((u, v, w))
    (p, k, a, b, c) = [int(word) for word in input().strip().split()]
    q = [int(word) for word in input().strip().split()]
    edge.sort(key=lambda e: e[-1])
    node = set([0, c])
    for i in range(m):
        node.add(edge[i][-1])
        for j in range(i + 1, m):
            w = edge[i][-1] + edge[j][-1]
            node.add((w + 1) // 2)
            node.add((w + 1) // 2 - 1)
    node = [x for x in node if x >= 0 and x <= c]
    arr = sorted(node)
    node = []
    mst = []
    for x in arr:
        res = kruskal(n, m, edge, x)
        if len(node) > 1:
            delta = mst[-1] - mst[-2]
            if delta % (node[-1] - node[-2]) == 0:
                delta = (mst[-1] - mst[-2]) // (node[-1] - node[-2])
                if mst[-1] + delta * (x - node[-1]) == res:
                    node.pop()
                    mst.pop()
        node.append(x)
        mst.append(res)
    ans = 0
    for i in range(p):
        ans ^= get_answer(node, mst, q[i])
    cur_q = q[-1]
    for i in range(p, k):
        cur_q = (cur_q * a + b) % c
        ans ^= get_answer(node, mst, cur_q)
    return ans

def main():
    print(solve())
if __name__ == '__main__':
    main()