Introduction

This is a challenging problem that combines segment tree merging with dynamic programming optimization. While it follows a relatively standard template, the overall difficulty level is high.

Prerequisites: Dynamic programming, tree-based DP, segmant tree with dynamic node allocation, segment tree merging.

Problem link: Click here

DP State Definition

The DP formulation for this problem has a unique structure that requires careful consideration.

First, sort all values and apply discretization. A special discretization technique is used here: without removing duplicates. This ensures (or at least guarantees) that every node has a distinct value.

Define f<sub>x,i</sub> as follows: Consider the subtree rooted at node x as an independent subproblem. For all 2<sup>sz(x)-1</sup> ways to cut edges within this subtree (sz(x) denotes the number of nodes in the subtree), count the total number of occurrences where the i-th smallest value appears as the minimum in a connected component that contains node x.

Consider an example where the value of node i equals i:

In this tree, treating the subtree rooted at node 2 as a new subproblem. For a particular edge-cutting configuration, suppose the connected component containing 2 is 3-2-5-6. The minimum value in this component is 2. This represents one valid edge-cutting scheme that contributes to f<sub>2,2</sub>.

The final answer can be expressed as:

∑(x ∈ V) [ 2^(n-sz(x)-1) × ∑(i=1 to n) f[x][i] × v[i] ]

Here, V represents the set of all nodes, and v[i] denotes the i-th value after special discretization. The term 2^(n-sz(x)-1) accounts for all possible configurations of edges outside the subtree, excluding the edge connecting the subtree root to its parent (which cannot be cut). Every such configuration can combine with any valid configuration inside the subtree.

This formulation covers all connected components exactly once.

DP Recurrence

The DP recurrence is non-trivial. For tree-based DP, consider the contribution when adding a child node y to node x. Two cases arise:

Case 1: Edge between x and y is kept.

For each pair (i, j), the value f[x][min(i,j)] receives a contribution of f[x][i] × f[y][j]. After simplification, the total contribution to f[x][i] becomes:

f[x][i] × ∑(j > i) f[y][j] + f[y][i] × ∑(j > i) f[x][j]

Note: There would normally be a term f[x][i] × f[y][i], but due to the special discretization method, both cannot be non-zero simultaneously, making this product always zero. This is precisely why the special discretization was employed.

Case 2: Edge between x and y is cut.

In this scenario, any edge-cutting configuration within y's subtree combines with any existing configuration in x's subtree to produce a valid answer. The resulting answer remains consistent with the case without y. Therefore, the contribution to f[x][i] is:

f[x][i] × 2^(sz(y)-1)

The term sz(y)-1 represents the number of edges within y's subtree, and 2^(sz(y)-1) counts the edge-cutting possibilities.

Combining both cases yields the final recurrrence:

f'[x][i] = f[x][i] × ( ∑(j > i) f[y][j] + 2^(sz(y)-1) ) + f[y][i] × ∑(j > i) f[x][j]

Data Structure Optimization

Examining the recurrence reveals suffix sum patterns. This suggests using segment tree merging to optimize the DP computation.

For each node x, maintain a dynamically allocated segment tree storing information about f[x]. Each node in the segment tree requires:

• sum: Sum of f[x][i] values in the interval, used for transitions
• dat: Sum of f[x][i] × value[i] in the interval, used for final answer computation

The segment tree must support lazy multiplication, so a multiplication lazy tag is also needed.

The core algorithm involves the segment tree merge operation, which efficiently combines f[y] into f[x].

The standard segment tree merge procedure recurses until one tree has an empty interval, or when reaching leaf nodes. Due to the special discretization, two leaf nodes will never both be non-empty simultaneously. When an interval in one tree is empty, all f[x][j] values are zero, making the suffix sum ∑(j > i) f[x][j] a constant value that has been maintained during recursion. Therefore, we simply multiply all values in the other tree's corresponding interval by this constant.

Implementation of the merge function (parameters p and q represent the segment trees for x and y respectively. Note that pmul includes the initial value 2^(sz(y)-1)):

int mergeTrees(int nodeP, int nodeQ, int left, int right, ll factorP, ll factorQ)
{
    if (nodeP == 0 && nodeQ == 0) return 0;
    if (nodeP == 0)
    {
        seg[nodeQ].applyMultiply(factorQ);
        return nodeQ;
    }
    if (nodeQ == 0)
    {
        seg[nodeP].applyMultiply(factorP);
        return nodeP;
    }
    propagate(nodeP);
    propagate(nodeQ);
    int mid = left + right >> 1;
    ll suffixP = seg[seg[nodeP].right].sum;
    ll suffixQ = seg[seg[nodeQ].right].sum;
    seg[nodeP].leftChild = mergeTrees(
        seg[nodeP].leftChild, seg[nodeQ].leftChild,
        left, mid,
        (factorP + suffixQ) % MOD, (factorQ + suffixP) % MOD
    );
    seg[nodeP].rightChild = mergeTrees(
        seg[nodeP].rightChild, seg[nodeQ].rightChild,
        mid + 1, right,
        factorP, factorQ
    );
    update(nodeP);
    return nodeP;
}

Complete Implementation

For space estimation: each call to Add traverses from root to leaf, creating at most ⌈log₂ n⌉ + 1 nodes (the height of the segment tree). Since this happens n times, the total space required is n × (⌈log₂ n⌉ + 1).

Full solution (approximately 4KB):

#include <bits/stdc++.h>

using namespace std;

namespace FastIO {
    const int BUFSIZE = 1 << 20;
    char buf[BUFSIZE], *p1 = buf, *p2 = buf;
    inline char getChar() {
        return (p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, BUFSIZE, stdin), p1 == p2))
            ? EOF : *p1++;
    }
    template<typename T>
    void readInt(T& x) {
        x = 0;
        bool neg = false;
        char ch = getChar();
        while (ch < '0' || ch > '9') {
            if (ch == '-') neg = true;
            ch = getChar();
        }
        while (ch >= '0' && ch <= '9') {
            x = x * 10 + (ch ^ '0');
            ch = getChar();
        }
        if (neg) x = -x;
    }
    template<typename T>
    void writeInt(T x) {
        if (!x) {
            putchar('0');
            return;
        }
        if (x < 0) {
            putchar('-');
            x = -x;
        }
        static int stk[55];
        int top = 0;
        while (x) {
            stk[++top] = x % 10;
            x /= 10;
        }
        while (top) putchar('0' + stk[top--]);
    }
}
using FastIO::readInt;
using FastIO::writeInt;

using ll = long long;

const int MAXN = 300005;
const int MAXLOG = 20;
const ll MOD = 998244353;

int vertexCount;
ll finalAnswer;

ll originalVals[MAXN], sortedVals[MAXN];
int discretized[MAXN], dupCount[MAXN];

void compressValues() {
    copy(originalVals + 1, originalVals + vertexCount + 1, sortedVals + 1);
    sort(sortedVals + 1, sortedVals + vertexCount + 1);
    for (int i = 1; i <= vertexCount; i++) {
        int pos = lower_bound(sortedVals + 1, sortedVals + vertexCount + 1, originalVals[i]) - sortedVals;
        discretized[i] = pos + dupCount[pos];
        dupCount[pos]++;
    }
}

struct Edge {
    int to, next;
} edges[MAXN << 1];
int head[MAXN], edgeCnt;

inline void addEdge(int u, int v) {
    edges[++edgeCnt] = {v, head[u]};
    head[u] = edgeCnt;
}

struct SegNode {
    int l, r;
    ll sum;
    ll lazyMult;
    ll weightedSum;
    int left, right;
    inline int middle() const { return l + r >> 1; }
    inline bool isLeaf() const { return l == r; }
    inline void multiply(ll v) {
        sum = sum * v % MOD;
        weightedSum = weightedSum * v % MOD;
        lazyMult = lazyMult * v % MOD;
    }
} segTree[MAXN * MAXLOG];
int segRoot[MAXN], nodePool;

inline void ensureNode(int& p, int L, int R) {
    if (!p) {
        p = ++nodePool;
        segTree[p] = {L, R, 0, 1, 0, 0, 0};
    }
}

inline void pullUp(int p) {
    const SegNode& leftNode = segTree[segTree[p].left];
    const SegNode& rightNode = segTree[segTree[p].right];
    segTree[p].sum = (leftNode.sum + rightNode.sum) % MOD;
    segTree[p].weightedSum = (leftNode.weightedSum + rightNode.weightedSum) % MOD;
}

inline void pushDown(int p) {
    ll mult = segTree[p].lazyMult;
    if (mult != 1) {
        int leftChild = segTree[p].left;
        int rightChild = segTree[p].right;
        if (leftChild) segTree[leftChild].multiply(mult);
        if (rightChild) segTree[rightChild].multiply(mult);
        segTree[p].lazyMult = 1;
    }
}

void insertValue(int pos, ll val, int& p) {
    if (segTree[p].isLeaf()) {
        segTree[p].sum = (segTree[p].sum + val) % MOD;
        segTree[p].weightedSum = (segTree[p].weightedSum + sortedVals[segTree[p].l] * val) % MOD;
        return;
    }
    int mid = segTree[p].middle();
    if (pos <= mid) {
        ensureNode(segTree[p].left, segTree[p].l, mid);
        insertValue(pos, val, segTree[p].left);
    } else {
        ensureNode(segTree[p].right, mid + 1, segTree[p].r);
        insertValue(pos, val, segTree[p].right);
    }
    pullUp(p);
}

int mergeTrees(int nodeP, int nodeQ, int L, int R, ll factorP, ll factorQ) {
    if (!nodeP && !nodeQ) return 0;
    if (!nodeP) {
        segTree[nodeQ].multiply(factorQ);
        return nodeQ;
    }
    if (!nodeQ) {
        segTree[nodeP].multiply(factorP);
        return nodeP;
    }
    pushDown(nodeP);
    pushDown(nodeQ);
    int mid = L + R >> 1;
    ll suffixP = segTree[segTree[nodeP].right].sum;
    ll suffixQ = segTree[segTree[nodeQ].right].sum;
    segTree[nodeP].left = mergeTrees(
        segTree[nodeP].left, segTree[nodeQ].left,
        L, mid,
        (factorP + suffixQ) % MOD, (factorQ + suffixP) % MOD
    );
    segTree[nodeP].right = mergeTrees(
        segTree[nodeP].right, segTree[nodeQ].right,
        mid + 1, R,
        factorP, factorQ
    );
    pullUp(nodeP);
    return nodeP;
}

ll power2[MAXN];
int subtreeSize[MAXN];

void dfs(int u, int parent) {
    subtreeSize[u] = 1;
    for (int i = head[u]; i; i = edges[i].next) {
        int v = edges[i].to;
        if (v == parent) continue;
        dfs(v, u);
        segRoot[u] = mergeTrees(segRoot[u], segRoot[v], 1, vertexCount,
                                power2[subtreeSize[v] - 1], 0);
        subtreeSize[u] += subtreeSize[v];
    }
    finalAnswer = (finalAnswer + power2[max(vertexCount - subtreeSize[u] - 1, 0)]
                   * segTree[segRoot[u]].weightedSum) % MOD;
}

int main() {
    readInt(vertexCount);
    for (int i = 1; i <= vertexCount; i++)
        readInt(originalVals[i]);
    for (int i = 1; i < vertexCount; i++) {
        int x, y;
        readInt(x);
        readInt(y);
        addEdge(x, y);
        addEdge(y, x);
    }
    compressValues();
    power2[0] = 1;
    for (int i = 1; i <= vertexCount; i++) {
        power2[i] = (power2[i - 1] << 1) % MOD;
        ensureNode(segRoot[i], 1, vertexCount);
        insertValue(discretized[i], 1, segRoot[i]);
    }
    dfs(1, 0);
    writeInt(finalAnswer);
    putchar('\n');
    return 0;
}