Vlad and the Mountains (Codeforces 1851G)
Tags: Implement, dsu

Vlad decided to go on a trip to the mountains. He plans to move between $$$n$$$ mountains, some of which are connected by roads. The $$$i$$$-th mountain has a height of $$$h_i$$$.If there is a road between mountains $$$i$$$ and $$$j$$$, Vlad can move from mountain $$$i$$$ to mountain $$$j$$$ by spending $$$h_j - h_i$$$ units of energy. If his energy drops below zero during the transition, he will not be able to move from mountain $$$i$$$ to mountain $$$j$$$. Note that $$$h_j - h_i$$$ can be negative and then the energy will be restored.Vlad wants to consider different route options, so he asks you to answer the following queries: is it possible to construct some route starting at mountain $$$a$$$ and ending at mountain $$$b$$$, given that he initially has $$$e$$$ units of energy?

Input: The first line of the input contains an integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases.The descriptions of the test cases follow.The first line of each test case contains two numbers $$$n$$$ and $$$m$$$ ($$$2 \le n \le 2 \cdot 10^5$$$, $$$1 \le m \le \min(\frac{n\cdot(n - 1)}{2}, 2 \cdot 10^5)$$$) — the number of mountains and the number of roads between them, respectively.The second line contains $$$n$$$ integers $$$h_1, h_2, h_3, \dots, h_n$$$ ($$$1 \le h_i \le 10^9$$$) — the heights of the mountains.The next $$$m$$$ lines contain two integers $$$u$$$ and $$$v$$$ ($$$1 \le u, v \le n$$$, $$$u \ne v$$$) — the numbers of the mountains connected by a road. It is guaranteed that no road appears twice.The next line contains an integer $$$q$$$ ($$$1 \le q \le 2 \cdot 10^5$$$) — the number of queries.The following $$$q$$$ lines contain three numbers $$$a$$$, $$$b$$$, and $$$e$$$ ($$$1 \le a, b \le n$$$, $$$0 \le e \le 10^9$$$) — the initial and final mountains of the route, and the amount of energy, respectively.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$. The same guarantee applies to $$$m$$$ and $$$q$$$.

Output: For each query, output "YES" if Vlad can construct a route from mountain $$$a$$$ to mountain $$$b$$$, and "NO" otherwise.You can output the answer in any case (for example, the strings "yEs", "yes", "Yes", and "YES" will be recognized as a positive answer).In the examples below, the answers for different test cases are separated by an empty line, which you do not need to output.


Sample Inputs:
2
7 7
1 5 3 4 2 4 1
1 4
4 3
3 6
3 2
2 5
5 6
5 7
5
1 1 3
6 2 0
4 7 0
1 7 4
1 7 2
6 5
4 7 6 2 5 1
1 3
5 3
1 5
2 4
6 2
5
1 5 1
1 3 1
1 2 1000
6 2 6
6 2 5

2
3 2
1 3 9
1 2
2 3
5
1 1 1
3 2 2
1 1 2
3 3 0
1 2 1
3 3
1 4 1
1 2
2 3
1 3
5
3 3 9
1 3 6
1 1 2
3 3 6
3 3 4

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

Sample Outputs:
YES
NO
YES
YES
NO

YES
NO
NO
YES
NO

YES
YES
YES
YES
NO

YES
YES
YES
YES
YES

YES
YES
YES
YES
YES

Solution (Python):
import sys
from collections import *
from math import sqrt
input = sys.stdin.buffer.readline

def comp(node):
    path = [node]
    while par[path[-1]] != path[-1]:
        path.append(par[path[-1]])
    for x in path:
        par[x] = path[-1]
    return path[-1]

def join(i, j):
    par[comp(i)] = comp(j)

def same(i, j):
    return comp(i) == comp(j)
for _ in range(int(input())):
    (n, m) = map(int, input().split())
    hs = [None] + list(map(int, input().split()))
    adj = defaultdict(list)
    for _ in range(m):
        edg = map(int, input().split())
        (lo, hi) = sorted(edg, key=lambda x: hs[x])
        adj[hi].append(lo)
    q = int(input())
    qs = []
    for qi in range(q):
        (fr, to, e) = map(int, input().split())
        mah = hs[fr] + e
        qs.append((mah, fr, to, qi))
    qs.sort()
    sorted_ms = sorted(range(1, n + 1), key=lambda x: hs[x])
    par = list(range(n + 1))
    ans = [None] * q
    next_m = 0
    for tu in qs:
        (mah, fr, to, qi) = tu
        while next_m < n and hs[sorted_ms[next_m]] <= mah:
            curr = sorted_ms[next_m]
            next_m += 1
            rds = adj[curr]
            for node in rds:
                join(node, curr)
        ans[qi] = 'YES' if same(fr, to) else 'NO'
    print(*ans, sep='\n')
    print()

Solution (Python):
import sys
from collections import *
from math import sqrt
input = sys.stdin.buffer.readline

def comp(node):
    path = [node]
    while par[path[-1]] != path[-1]:
        path.append(par[path[-1]])
    for x in path:
        par[x] = path[-1]
    return path[-1]

def join(i, j):
    par[comp(i)] = comp(j)

def same(i, j):
    return comp(i) == comp(j)
for _ in range(int(input())):
    (n, m) = map(int, input().split())
    hs = [None] + list(map(int, input().split()))
    adj = defaultdict(list)
    for _ in range(m):
        edg = map(int, input().split())
        (lo, hi) = sorted(edg, key=lambda x: hs[x])
        adj[hi].append(lo)
    q = int(input())
    qs = []
    for qi in range(q):
        (fr, to, e) = map(int, input().split())
        mah = hs[fr] + e
        qs.append((mah, fr, to, qi))
    qs.sort()
    sorted_ms = sorted(range(1, n + 1), key=lambda x: hs[x])
    par = list(range(n + 1))
    ans = [None] * q
    next_m = 0
    for tu in qs:
        (mah, fr, to, qi) = tu
        while next_m < n and hs[sorted_ms[next_m]] <= mah:
            curr = sorted_ms[next_m]
            next_m += 1
            rds = adj[curr]
            for node in rds:
                join(node, curr)
        ans[qi] = 'YES' if same(fr, to) else 'NO'
    print(*ans, sep='\n')
    print()

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 main():
    t = int(input())
    allans = []
    for _ in range(t):
        (n, m) = readIntArr()
        h = [-1] + readIntArr()
        edges = []
        for __ in range(m):
            (u, v) = readIntArr()
            edges.append((u, v))
        q = int(input())
        queries = []
        for __ in range(q):
            (a, b, e) = readIntArr()
            queries.append((a, b, e))
        edges_height_order = list(range(m))
        edges_height_order.sort(key=lambda i: max(h[edges[i][0]], h[edges[i][1]]))
        ans = [-1] * q
        queries_upper_height_order = list(range(q))
        queries_upper_height_order.sort(key=lambda i: h[queries[i][0]] + queries[i][2])
        uf = UnionFind(n + 1)
        i = 0
        for query_i in queries_upper_height_order:
            (a, b, e) = queries[query_i]
            while i < m:
                j = edges_height_order[i]
                (x, y) = edges[j]
                if max(h[x], h[y]) <= h[a] + e:
                    uf.union(x, y)
                    i += 1
                else:
                    break
            if uf.find(a) == uf.find(b):
                ans[query_i] = 'YES'
            else:
                ans[query_i] = 'NO'
        allans.extend(ans)
        allans.append('')
    multiLineArrayPrint(allans)
    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):
    print('? {} {}'.format(a, b))
    sys.stdout.flush()
    return int(input())

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