Imbalance Value of a Tree (Codeforces 915F)
Tags: Others, dsu

You are given a tree T consisting of n vertices. A number is written on each vertex; the number written on vertex i is ai. Let's denote the function I(x, y) as the difference between maximum and minimum value of ai on a simple path connecting vertices x and y.Your task is to calculate .

Input: The first line contains one integer number n (1 ≤ n ≤ 106) — the number of vertices in the tree.The second line contains n integer numbers a1, a2, ..., an (1 ≤ ai ≤ 106) — the numbers written on the vertices.Then n - 1 lines follow. Each line contains two integers x and y denoting an edge connecting vertex x and vertex y (1 ≤ x, y ≤ n, x ≠ y). It is guaranteed that these edges denote a tree.

Output: Print one number equal to .


Sample Inputs:
4
2 2 3 1
1 2
1 3
1 4

Sample Outputs:
6

Solution (Python):
from array import array
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def get_root(s):
    v = []
    while not s == root[s]:
        v.append(s)
        s = root[s]
    for i in v:
        root[i] = s
    return s

def unite(s, t):
    (rs, rt) = (get_root(s), get_root(t))
    if not rs ^ rt:
        return 0
    if get_size(rs) >= get_size(rt):
        root[rt] = rs
        size[rs] += size[rt]
        return rs
    else:
        root[rs] = rt
        size[rt] += size[rs]
        return rt

def same(s, t):
    return True if get_root(s) == get_root(t) else False

def get_size(s):
    return size[get_root(s)]
n = int(input())
a = list(map(int, input().split()))
G = [array('i', []) for _ in range(n + 1)]
for _ in range(n - 1):
    (x, y) = map(int, input().split())
    G[x].append(y)
    G[y].append(x)
l = max(a)
b = [0 for _ in range(l + 1)]
for i in range(1, n + 1):
    if not b[a[i - 1]]:
        b[a[i - 1]] = [i]
    else:
        b[a[i - 1]].append(i)
root = array('i', [i for i in range(n + 1)])
rank = array('i', [1 for _ in range(n + 1)])
size = array('i', [1 for _ in range(n + 1)])
(ans1, ans2) = (0, 0)
mod = pow(10, 9) + 7
for i in range(1, l + 1):
    if not b[i]:
        continue
    ans0 = 0
    for j in b[i]:
        s = get_size(j)
        aj = a[j - 1]
        for k in G[j]:
            if a[k - 1] < aj or (a[k - 1] == aj and k < j):
                x = get_size(k)
                ans0 += s * x
                s += x
                unite(j, k)
    ans1 += ans0 * i % mod
    ans2 += ans1 // mod
    ans1 %= mod
    ans2 += ans0 * i // mod
for i in range(1, n + 1):
    root[i] = i
    rank[i] = 1
    size[i] = 1
for i in range(l, 0, -1):
    if not b[i]:
        continue
    ans0 = 0
    for j in b[i]:
        s = get_size(j)
        aj = a[j - 1]
        for k in G[j]:
            if a[k - 1] > aj or (a[k - 1] == aj and k < j):
                x = get_size(k)
                ans0 += s * x
                s += x
                unite(j, k)
    ans1 -= ans0 * i % mod
    if ans1 < 0:
        ans2 -= 1
        ans1 %= mod
    ans2 -= ans0 * i // mod
ans = ans1 + ans2 * mod
print(ans)

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def get_root(s):
    v = []
    while not s == root[s]:
        v.append(s)
        s = root[s]
    for i in v:
        root[i] = s
    return s

def unite(s, t):
    (rs, rt) = (get_root(s), get_root(t))
    if not rs ^ rt:
        return 0
    if get_size(rs) >= get_size(rt):
        root[rt] = rs
        size[rs] += size[rt]
        return rs
    else:
        root[rs] = rt
        size[rt] += size[rs]
        return rt

def same(s, t):
    return True if get_root(s) == get_root(t) else False

def get_size(s):
    return size[get_root(s)]

def make_graph(n, m):
    G = []
    for _ in range(m):
        (u, v) = map(int, input().split())
        G.append(u * (n + 1) + v)
        G.append(v * (n + 1) + u)
    G.sort()
    m *= 2
    s = [m] * (n + 2)
    for i in range(m):
        u = G[i] // (n + 1)
        if s[u] == m:
            s[u] = i
        G[i] %= n + 1
    return (G, s)
n = int(input())
a = list(map(int, input().split()))
(G, u) = make_graph(n, n - 1)
b = [a[i - 1] * (n + 1) + i for i in range(1, n + 1)]
b.sort()
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
ans = 0
for i in b:
    j = i % (n + 1)
    ans0 = 0
    s = get_size(j)
    aj = a[j - 1]
    for m in range(u[j], u[j + 1]):
        k = G[m]
        if a[k - 1] < aj or (a[k - 1] == aj and k < j):
            x = get_size(k)
            ans0 += s * x
            s += x
            unite(j, k)
    ans += i // (n + 1) * ans0
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
for i in b[::-1]:
    j = i % (n + 1)
    ans0 = 0
    s = get_size(j)
    aj = a[j - 1]
    for m in range(u[j], u[j + 1]):
        k = G[m]
        if a[k - 1] > aj or (a[k - 1] == aj and k < j):
            x = get_size(k)
            ans0 += s * x
            s += x
            unite(j, k)
    ans -= i // (n + 1) * ans0
print(ans)

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def get_root(s):
    v = []
    while not s == root[s]:
        v.append(s)
        s = root[s]
    for i in v:
        root[i] = s
    return s

def unite(s, t):
    (rs, rt) = (get_root(s), get_root(t))
    if not rs ^ rt:
        return 0
    if get_size(rs) >= get_size(rt):
        root[rt] = rs
        size[rs] += size[rt]
        return rs
    else:
        root[rs] = rt
        size[rt] += size[rs]
        return rt

def same(s, t):
    return True if get_root(s) == get_root(t) else False

def get_size(s):
    return size[get_root(s)]

def make_graph(n, m):
    G = []
    for _ in range(m):
        (u, v) = map(int, input().split())
        G.append(u * (n + 1) + v)
        G.append(v * (n + 1) + u)
    G.sort()
    m *= 2
    s = [m] * (n + 2)
    for i in range(m):
        u = G[i] // (n + 1)
        if s[u] == m:
            s[u] = i
        G[i] %= n + 1
    return (G, s)
n = int(input())
a = list(map(int, input().split()))
G = [[] for _ in range(n + 1)]
(G, u) = make_graph(n, n - 1)
b = []
for i in range(1, n + 1):
    b.append(a[i - 1] * (n + 1) + i)
b.sort()
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
ans = 0
for i in b:
    j = i % (n + 1)
    ans0 = 0
    s = get_size(j)
    aj = a[j - 1]
    for m in range(u[j], u[j + 1]):
        k = G[m]
        if a[k - 1] < aj or (a[k - 1] == aj and k < j):
            x = get_size(k)
            ans0 += s * x
            s += x
            unite(j, k)
    ans += i // (n + 1) * ans0
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
for i in b[::-1]:
    j = i % (n + 1)
    ans0 = 0
    s = get_size(j)
    aj = a[j - 1]
    for m in range(u[j], u[j + 1]):
        k = G[m]
        if a[k - 1] > aj or (a[k - 1] == aj and k < j):
            x = get_size(k)
            ans0 += s * x
            s += x
            unite(j, k)
    ans -= i // (n + 1) * ans0
print(ans)