Problem A: Graph Coloring

We are given an integer n and need to color the integers from 1 to n. The constraint is that for any two integers i and j where i < j, if their difference j - i is a prime number, they must have differant colors. The goal is to use the minimum number of colors possible and provide a valid coloring scheme.

For n >= 7, it can be proven that the minimum number of colors required is 4. A valid coloring scheme is to repeat the sequence 1, 2, 3, 4. In this pattern, any two numbers with the same color will have a difference that is a multiple of 4. Since 4 is the first composite number, any multiple of 4 greater than 4 itself cannot be a prime number, thus satisfying the constraint. For smaller values of n (from 1 to 6), we can determine the optimal number of colors by manual inspection.

#include <bits>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;

    // Precomputed optimal color counts for small n
    int color_counts[] = {1, 1, 2, 2, 3, 3};
    string colorings[] = {
        "1",
        "1 1",
        "1 1 2",
        "1 1 2 2",
        "1 1 2 2 3",
        "1 1 2 2 3 3"
    };

    if (n <= 6) {
        cout << color_counts[n - 1] << "
" << colorings[n - 1] << "
";
    } else {
        cout << "4
";
        for (int i = 1; i <= n; ++i) {
            cout << (i - 1) % 4 + 1 << " ";
        }
        cout << "
";
    }

    return 0;
}
</bits>

Problem B: Sequence Construction

Given a sequence of integers a of length m, an integer n, and a budget D, we need to construct a sequence b of the same length. Each element b_i must be in the range [0, n], and the sum of products sum(a_i * b_i) must not exceed D. The objective is to maximize the value of sum(b_i) + k * min(b_i), where k is a given constant.

The objective function, when viewed as a function of the minimum value mn in the sequence b, is a unimodal function (it first increases, reaches a peak, and then decreases). This property allows us to use ternary search to efficiently find the optimal value for mn that maximizes the objective.

#include <bits>
using namespace std;

using ll = long long;

ll n, k, D;
int m;
vector<ll> a;
vector<ll> prefix_sum;

ll calculate_objective(ll mn) {
    if (mn < 0) return 0;
    ll total = mn * k + mn * m;
    ll remaining_budget = D - mn * prefix_sum.back();

    if (remaining_budget < 0) return 0;

    // Find the largest index where (n - mn) * a[i] <= remaining_budget
    int l = 0, r = m - 1, pos = 0;
    while (l <= r) {
        int mid = l + (r - l) / 2;
        if ((n - mn) * a[mid] <= remaining_budget) {
            pos = mid;
            l = mid + 1;
        } else {
            r = mid - 1;
        }
    }

    total += (n - mn) * pos;
    remaining_budget -= (n - mn) * a[pos];

    // Use remaining budget for the next element
    if (pos + 1 < m) {
        ll additional = min(n - mn, remaining_budget / a[pos + 1]);
        total += additional;
    }

    return total;
}

ll ternary_search(ll low, ll high) {
    while (high - low > 2) {
        ll mid1 = low + (high - low) / 3;
        ll mid2 = high - (high - low) / 3;

        ll val1 = calculate_objective(mid1);
        ll val2 = calculate_objective(mid2);

        if (val1 >= val2) {
            high = mid2;
        } else {
            low = mid1;
        }
    }
    // Check a few points around the final range for safety
    ll best = 0;
    for (ll x = low; x <= high; ++x) {
        best = max(best, calculate_objective(x));
    }
    return best;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int T;
    cin >> T;
    while (T--) {
        cin >> n >> m >> k >> D;
        a.resize(m);
        prefix_sum.resize(m + 1, 0);

        for (int i = 0; i < m; ++i) {
            cin >> a[i];
        }
        sort(a.begin(), a.end());
        for (int i = 0; i < m; ++i) {
            prefix_sum[i + 1] = prefix_sum[i] + a[i];
        }

        ll optimal_mn = ternary_search(0, min(n, D / prefix_sum.back()));
        cout << calculate_objective(optimal_mn) << "
";
    }

    return 0;
}
</ll></ll></bits>

Problem C: Path Queries on a Tree

Given a tree with n nodes, we need to answer m queries. Each query provides two nodes l and r, and we must count the number of nodes k on the path from l to r such that the number of steps taken to reach k from l is exactly equal to the value of k.

The path from l to r can be split into two segments: from l to the Lowest Common Ancestor (LCA) and from the LCA to r. For the first segment, the condition dist(l, k) = k translates to dep[l] = dep[k] + k. For the second segment, the condition becomes dep[k] - k = 2 * dep[lca] - dep[l]. By defining appropriate offsets, we can use tree decomposition (heavy-light decomposition) combined with a Persistent Segment Tree to preprocess the tree and answer each query in logarithmic time.

#include <bits>
using namespace std;

const int MAXN = 3e5 + 10;
const int OFFSET = 3e5;

int n, m;
vector<int> adj[MAXN];
int depth[MAXN], parent[MAXN], heavy[MAXN], head[MAXN], pos[MAXN], timer;

struct PersistentSegmentTree {
    struct Node {
        int left, right, count;
    };
    vector<node> nodes;
    vector<int> roots;

    PersistentSegmentTree(int size) : nodes(1), roots(size + 1) {}

    int new_node() { nodes.push_back({0, 0, 0}); return nodes.size() - 1; }

    void build(int v, int tl, int tr) {
        if (tl == tr) return;
        int tm = (tl + tr) / 2;
        nodes[v].left = new_node();
        nodes[v].right = new_node();
        build(nodes[v].left, tl, tm);
        build(nodes[v].right, tm + 1, tr);
    }

    void init(int size) {
        roots[0] = new_node();
        build(0, 1, size);
    }

    int update(int v, int tl, int tr, int idx) {
        int u = new_node();
        nodes[u] = nodes[v];
        if (tl == tr) {
            nodes[u].count++;
            return u;
        }
        int tm = (tl + tr) / 2;
        if (idx <= tm) {
            nodes[u].left = update(nodes[v].left, tl, tm, idx);
        } else {
            nodes[u].right = update(nodes[v].right, tm + 1, tr, idx);
        }
        nodes[u].count = nodes[nodes[u].left].count + nodes[nodes[u].right].count;
        return u;
    }

    int query(int v, int u, int tl, int tr, int l, int r) {
        if (r < tl || tr < l) return 0;
        if (l <= tl && tr <= r) return nodes[u].count - nodes[v].count;
        int tm = (tl + tr) / 2;
        return query(nodes[v].left, nodes[u].left, tl, tm, l, r) +
               query(nodes[v].right, nodes[u].right, tm + 1, tr, l, r);
    }
};

void dfs(int v, int p = -1) {
    parent[v] = p;
    heavy[v] = -1;
    int max_size = 0;
    for (int u : adj[v]) {
        if (u == p) continue;
        depth[u] = depth[v] + 1;
        dfs(u, v);
        if (max_size < adj[u].size()) {
            max_size = adj[u].size();
            heavy[v] = u;
        }
    }
}

void decompose(int v, int h) {
    head[v] = h;
    pos[v] = ++timer;
    if (heavy[v] != -1) {
        decompose(heavy[v], h);
    }
    for (int u : adj[v]) {
        if (u != parent[v] && u != heavy[v]) {
            decompose(u, u);
        }
    }
}

int lca(int u, int v) {
    while (head[u] != head[v]) {
        if (depth[head[u]] > depth[head[v]]) swap(u, v);
        v = parent[head[v]];
    }
    return depth[u] < depth[v] ? u : v;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n >> m;
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    depth[1] = 1;
    dfs(1);
    decompose(1, 1);

    PersistentSegmentTree pst(2 * n + OFFSET);
    pst.init(2 * n + OFFSET);

    for (int i = 1; i <= n; ++i) {
        int val1 = pos[i] + depth[i] + OFFSET;
        pst.roots[i] = pst.update(pst.roots[i - 1], 1, 2 * n + OFFSET, val1);
    }

    for (int i = 1; i <= m; ++i) {
        int l, r;
        cin >> l >> r;
        int ancestor = lca(l, r);
        int ans = 0;

        // Query path l to lca
        while (head[l] != head[ancestor]) {
            ans += pst.query(pst.roots[pos[head[l]] - 1], pst.roots[pos[l]], 1, 2 * n + OFFSET, depth[l] + OFFSET, 2 * n + OFFSET);
            l = parent[head[l]];
        }
        ans += pst.query(pst.roots[pos[ancestor] - 1], pst.roots[pos[l]], 1, 2 * n + OFFSET, depth[l] + OFFSET, 2 * n + OFFSET);

        // Query path r to lca
        pst = PersistentSegmentTree(2 * n + OFFSET); // Re-initialize for second query
        pst.init(2 * n + OFFSET);
        for (int j = 1; j <= n; ++j) {
            int val2 = depth[j] - pos[j] + OFFSET;
            pst.roots[j] = pst.update(pst.roots[j - 1], 1, 2 * n + OFFSET, val2);
        }

        while (head[r] != head[ancestor]) {
            ans += pst.query(pst.roots[pos[head[r]] - 1], pst.roots[pos[r]], 1, 2 * n + OFFSET, 1, 2 * depth[ancestor] - depth[l] + OFFSET);
            r = parent[head[r]];
        }
        ans += pst.query(pst.roots[pos[ancestor] - 1], pst.roots[pos[r]], 1, 2 * n + OFFSET, 1, 2 * depth[ancestor] - depth[l] + OFFSET);

        // Adjust for the LCA node if it was counted twice
        if (ancestor + depth[ancestor] == depth[l]) ans--;
        
        cout << ans << "
";
    }

    return 0;
}
</int></node></int></bits>

Problem D: Counting Unique Subarrays

We have a sequence of integers. The sequence supports two operations: a single-point update (changing the value of one element) and a query. A query asks for the number of subarrays within a given range [l, r] that contain no duplicate elements.

To solve this, we can preprocess an array max_right[i] where max_right[i] represents the rightmost index of the longest subarray starting at i that contains all unique elements. This array can be maintained efficiently using a two-pointer sliding window technique. For each update, we recompute the affected segments of max_right. For a query [l, r], the answer is the sum of min(max_right[i], r) - i + 1 for all i from l to r.

#include <bits>
using namespace std;

const int MAXN = 2e5 + 10;

int n, m;
int a[MAXN];
int max_right[MAXN];

void compute_max_right() {
    vector<bool> seen(n + 2, false);
    int left_ptr = 1, right_ptr = 0;
    while (left_ptr <= n) {
        while (right_ptr < n && !seen[a[right_ptr + 1]]) {
            seen[a[++right_ptr]] = true;
        }
        max_right[left_ptr] = right_ptr;
        seen[a[left_ptr++]] = false;
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n >> m;
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
    }
    compute_max_right();

    while (m--) {
        int op, l, r;
        cin >> op >> l >> r;
        if (op == 1) {
            a[l] = r;
            compute_max_right();
        } else if (op == 2) {
            long long ans = 0;
            for (int i = l; i <= r; ++i) {
                ans += min(max_right[i], r) - i + 1;
            }
            cout << ans << "
";
        }
    }

    return 0;
}
</bool></bits>