To begin with, we note that it makes no sense to buy the same number of pumpkins more than 3 times, and it also makes no sense to buy more than $$$(3^{\lfloor \log_3 n \rfloor})$$$ pumpkins at once. Indeed, if we buy $$$(3^k)$$$ pumpkins 3 times, we will end up with $$$(3^{k+1})$$$ pumpkins , and we could have bought $$$(3^{k+1})$$$ pumpkins in one deal, which would have resulted in 2 fewer deals. The second statement means that it makes no sense to buy more pumpkins than necessary. Thus, we can assign an integer from 0 to 2 to all possible deal options. Let us assign the number $$$(w_i)$$$ to the $$$(i)$$$-th deal option; then the following condition must be satisfied:

As can be observed, this is the definition of the representation of the number $$$(n)$$$ in the ternary numeral system. The representation in the ternary system is unique, which means we will find this representation $$$((w_0,w_1,\ldots,w_{\lfloor \log_3 n \rfloor}))$$$, and the answer, as can be seen, will be:

#include <bits/stdc++.h>
using namespace std;
 
int main() {
    vector <long long> cost;
    long long c = 3;
    long long cnt = 1;
    for (int i = 0; i < 21; ++i) {
        cost.push_back(c);
        c = 3 * c + cnt;
        cnt *= 3;
    }

    int t;
    cin >> t;
    while (t--) {
        long long n;
        cin >> n;
        long long min_k = 0;
        long long min_cost = 0;
        int sz = 0;
        while (n) {
            min_k += n % 3;
            min_cost += (n % 3) * cost[sz];
            n /= 3;
            sz++;
        }
        cout << min_cost << '\n';
    }
    
    return 0;
}

There are two solutions:

1. We can make partial sums $$$(sum_i = a_1 + a_2 + \ldots + a_i)$$$ and then make a binary search for each query $$$q_i$$$ to find the result $$$j$$$ with the properties $$$sum_{j-1} \lt q_i$$$ and

   Unable to parse markup [type=CF_MATHJAX]

   .

2. We can precalculate the index of the pile for each worm and then answer for each query in $$$O(1)$$$. This solution has the complexity $$$O(n + m)$$$.

#include <bits/stdc++.h>
using namespace std;

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

    int n;
    cin >> n;

    vector<long long> pref(n);
    for (int i = 0; i < n; i++) {
        cin >> pref[i];
        if (i > 0) pref[i] += pref[i - 1];
    }

    int m;
    cin >> m;

    while (m--) {
        long long q;
        cin >> q;

        int ans = lower_bound(pref.begin(), pref.end(), q)
                  - pref.begin();

        cout << ans + 1 << '\n';
    }

    return 0;
}

#include <bits/stdc++.h>
using namespace std;

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

    int n;
    cin >> n;

    vector<long long> a(n);
    long long total = 0;

    for (int i = 0; i < n; i++) {
        cin >> a[i];
        total += a[i];
    }

    vector<int> bag(total + 1);

    ll label = 1;
    for (int i = 0; i < n; i++) {
        for (ll j = 0; j < a[i]; j++) {
            bag[label++] = i + 1;
        }
    }

    int m;
    cin >> m;

    while (m--) {
        ll q;
        cin >> q;
        cout << bag[q] << '\n';
    }

    return 0;
}

This problem with small coordinates can be solved using partial sums and some easy counting. Let's carry an array $$$cnt$$$, where $$$cnt_i$$$ will be equal to the number of segments that cover the point with coordinate $$$i$$$. How to calculate $$$cnt$$$ in $$$(O(n+\max X)$$$?

For each segment $$$(l_i,r_i)$$$, let's add $$$+1$$$ to $$$cnt_{l_i}$$$ and $$$-1$$$ to $$$cnt_{r_i+1}$$$. Now build prefix sums on this array and notice that $$$cnt_i$$$ equals the number of segments that cover the point with coordinate $$$i$$$. Then $$$ans_i$$$ will be equal to

All the answers can be calculated in $$$O(\max X)$$$ in total. So the total complexity of this solution is $$$O(n+\max X)$$$.

But in our problem it is too slow to build an entire array $$$cnt$$$. So what should we do? It is obvious that if any coordinate $$$j$$$ is not equal to some $$$l_i$$$ or some $$$r_i+1$$$, then $$$cnt_j=cnt_{j-1}$$$. So we do not need to carry all the positions explicitly. Let's carry all $$$l_i$$$ and $$$r_i+1$$$ in some logarithmic data structure or let's use the coordinate compression method.

The coordinate compression method allows us to transform the set of big sparse objects to the set of small compressed objects while maintaining the relative order. In our problem, let's do the following things: push all $$$l_i$$$ and $$$r_i+1$$$ into vector $$$cval$$$, sort this vector, keep only unique values, and then use the position of elements in vector $$$cval$$$ instead of the original value (any position can be found in $$$O(\log n)$$$ by binary search or standard methods such as lower_bound in C++).

So the first part of the solution works in $$$O(n\log n)$$$. The answer can be calculated using almost the same approach as in the solution to this problem with small coordinates. But now we know that between two adjacent elements $$$cval_i$$$ and $$$cval_{i+1}$$$, there are exactly $$$cval_{i+1}-cval_i$$$ points with answer equal to $$$cnt_i$$$. So if we iterate over all pairs of adjacent elements $$$cval_i$$$ and $$$cval_{i+1}$$$ and add $$$cval_{i+1}-cval_i$$$ to $$$ans_{cnt_i}$$$, we will calculate all the answers in $$$O(n)$$$.

Time complexity: $$$O(n log(n))$$$

#include <bits/stdc++.h>
using namespace std;

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

    int n;
    cin >> n;

    vector<pair<long long, long long>> segments(n);
    vector<long long> coords;

    for (auto &[l, r] : segments) {
        cin >> l >> r;
        coords.push_back(l);
        coords.push_back(r + 1);
    }

    sort(coords.begin(), coords.end());
    coords.erase(unique(coords.begin(), coords.end()), coords.end());

    int k = coords.size();
    vector<int> cnt(k, 0);

    for (auto [l, r] : segments) {
        int left = lower_bound(coords.begin(), coords.end(), l)
                   - coords.begin();
        int right = lower_bound(coords.begin(), coords.end(), r + 1)
                    - coords.begin();

        cnt[left]++;
        cnt[right]--;
    }

    for (int i = 1; i < k; i++)
        cnt[i] += cnt[i - 1];

    vector<long long> ans(n + 1, 0);

    for (int i = 1; i < k; i++)
        ans[cnt[i - 1]] += coords[i] - coords[i - 1];

    for (int i = 1; i <= n; i++)
        cout << ans[i] << (i == n ? '\n' : ' ');

    return 0;
}

Let's look at the first nail. If it is occupied by the fold place, then the Witch will put the next fold place on the third nail, then on the fifth, and so on. Otherwise, if the first nail is occupied by an end of a rod, then the second, fourth, sixth, and so on nails will be occupied by the fold places.

Let's see if we can complete our polyline with the first nail occupied by a fold place. This means we should check whether we have an unused rod with length $$$\operatorname{dist}(\text{nails}[n],\text{nails}[1])+\operatorname{dist}(\text{nails}[1],\text{nails}[2])$$$. Then check the third nail, and so on.

If we can complete the polyline, we have found a valid answer. Otherwise, repeat the previous procedure, but start from the second nail, assuming that the first nail is occupied by an end of a rod instead.

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n, m;
    cin >> n >> m;

    vector<long long> x(n), y(n), d(n), r(m);

    for (int i = 0; i < n; i++)
        cin >> x[i] >> y[i];

    for (int i = 0; i < m; i++)
        cin >> r[i];

    for (int i = 0; i < n; i++) {
        int j = (i + 1) % n;
        d[i] = abs(x[i] - x[j]) + abs(y[i] - y[j]);
    }

    for (int s = 0; s < 2; s++) {
        map<long long, vector<int>> mp;

        for (int i = 0; i < m; i++)
            mp[r[i]].push_back(i + 1);

        vector<int> ans(n, -1);
        bool ok = true;

        for (int i = s; i < n; i += 2) {
            long long need = d[(i - 1 + n) % n] + d[i];

            if (mp[need].empty()) {
                ok = false;
                break;
            }

            ans[i] = mp[need].back();
            mp[need].pop_back();
        }

        if (ok) {
            cout << "YES\n";

            for (int i = 0; i < n; i++)
                cout << ans[i] << " \n"[i == n - 1];

            return 0;
        }
    }

    cout << "NO\n";
    return 0;
}

Solution 1: Dynamic Programming on Subtrees

We define dp[u][i] as a boolean indicating whether it is possible to collect all candies in the subtree of vertex u while shaking a number of vertices in that subtree congruent to i modulo 3.

Base Case

If vertex u is a leaf, we must shake it to collect its candy. Therefore,

• dp[u][0] = false
• dp[u][1] = true
• dp[u][2] = false

Transition

For a non-leaf vertex u, we consider two cases.

Case 1: Shake vertex u.

If we shake u, all candies in its subtree fall at once. We cannot shake any of its descendants, so the number of shaken vertices in this subtree is exactly one.

Therefore, dp[u][1] = true.

Case 2: Do not shake vertex u.

In this case, we must collect all candies by shaking vertices inside the subtrees of its children.

We maintain an auxiliary array dpson[j], which indicates whether it is possible to collect all candies in the subtrees of the children processed so far, with the total number of shaken vertices congruent to j modulo 3.

Initially,

• dpson[0] = true
• dpson[1] = false
• dpson[2] = false

For every child v of u, we combine the states of dpson and dp[v].

If we shake j vertices in the already processed subtrees and i vertices in the subtree of v, the total number of shaken vertices becomes (i + j) % 3.

The transition is:

for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 3; j++) {
        if (dp[v][i] && dpson[j]) {
            newdp[(i + j) % 3] = true;
        }
    }
}

After processing child v, we replace dpson with newdp and continue with the next child.

Once all children have been processed, we assign dpson to dp[u]. Since shaking u itself is also a valid option, we set dp[u][1] = true.

We process the tree using DFS in postorder so that the states of all children are calculated before their parent.

Final Answer

The root is vertex 1. We need to collect all candies while shaking a number of vertices divisible by 3.

Therefore, the answer is YES if dp[1][0] is true; otherwise, it is NO.

Complexity Analysis

• Time Complexity: $$$O(n)$$$, since each edge is processed a constant number of times.
• Space Complexity: $$$O(n)$$$, for the tree and the DP states.

Solution 2: Analytical Approach

Instead of calculating DP states, we can derive a simpler condition that determines the answer.

Key Observation

Initially, consider the set containing only the root. We can repeatedly perform the following operation:

Choose a non-leaf vertex u in the set and replace it with all its children.

This operation preserves the validity of the set: the selected vertices remain pairwise unrelated by ancestry, and every leaf is eventually represented by exactly one selected vertex.

Thus, every valid set of vertices that collects all candies can be obtained through these operations.

Initially, the set contains one vertex. Replacing a vertex with its children changes the size of the set by

$ \Delta = \text{children}(u)-1 $

Therefore, only the value of this change modulo 3 matters.

Define

where $$$c_u$$$ is the number of children of vertex u.

Replacing vertex u with its children changes the size of the selected set by $$$a[u]$$$ modulo 3.

Vertices for which a[u] = 0 do not affect the size modulo 3, so we only need to consider vertices for which a[u] is nonzero.

When Is the Answer NO?

Suppose there are two vertices u and v such that neither is an ancestor of the other, and both have nonzero values.

We can perform operations on all other vertices first, without expanding u or v. Let the current size of the selected set modulo 3 be S.

• If S = 0, we have already found a valid solution.
• If either a[u] or a[v] equals 3 - S, expanding that vertex makes the size divisible by 3.
• Otherwise, both values must equal S. Expanding both vertices changes the size by S + S, making the final size congruent to 0 modulo 3.

Hence, if two nonzero-valued vertices are not on the same root-to-leaf path, the answer is always YES.

Therefore, for the answer to be NO, all vertices with nonzero values must lie on a single path starting at the root.

There is one more condition. Let these vertices appear in order from the root downwards.

• The first vertex must have value 1.
• The second must have value 2.
• The third must have value 1.
• The fourth must have value 2, and so on.

Why? If the first vertex has value 2, expanding it changes the initial set size from 1 to 0 modulo 3, giving a valid solution immediately. Similarly, at every subsequent vertex, its value must be the opposite of the previous nonzero value; otherwise, expanding the appropriate vertices makes the size divisible by 3.

Thus, the answer is NO if and only if all nonzero-valued vertices lie on a single root-to-leaf path and their values alternate as

$ 1,2,1,2,\ldots $

In every other case, the answer is YES.

Algorithm

We can check this condition using a single DFS:

1. Root the tree at vertex 1.
2. Calculate the number of children of every vertex.
3. Calculate a[u] for each vertex.
4. Check whether all vertices with nonzero values lie on one root-to-leaf path.
5. Check whether their values alternate between 1 and 2, starting with 1.
6. If both conditions hold, output NO; otherwise, output YES.

Complexity Analysis

• Time Complexity: (O(n)), because each vertex and edge is processed a constant number of times.
• Space Complexity: (O(n)), for storing the tree and auxiliary information.

The analytical approach is simpler to implement and avoids maintaining the three DP states for every vertex.

#include <bits/stdc++.h>
 
using namespace std;
 
const int MAXN = 200001;

int n;
vector <int> graph[MAXN];
vector <bool> dp[MAXN];

void merge_dp(vector <bool>& lhs, vector <bool>& rhs, vector <bool>& result) {
    for (int i = 0; i < 3; ++i) {
        if (!lhs[i]) continue;
        for (int j = 0; j < 3; ++j) {
            if (!rhs[j]) continue;
            result[(i + j) % 3] = true;
        }
    }
}

void dfs(int u, int p) {
    vector <int> child;

    for (auto v : graph[u]) if (v != p) {
        dfs(v, u);
        child.push_back(v);
    }

    dp[u].resize(3);

    if (!child.empty()) {
        dp[u][0] = true;
        for (auto v : child) {
            vector <bool> result(3, false);
 
            merge_dp(dp[u], dp[v], result);
            dp[u] = result;
        }
    }

    dp[u][1] = true;
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int t;
    cin >> t;
    while (t--) {
        cin >> n;
        for (int i = 1; i <= n; ++i) {
            graph[i].clear();
            dp[i].clear();
        }
        for (int i = 0; i < n - 1; ++i) {
            int u, v;
            cin >> u >> v;
            graph[u].push_back(v);
            graph[v].push_back(u);
        }

        dfs(1, -1);

        cout << (dp[1][0] ? "YES" : "NO") << '\n';
    }

    return 0;
}

#include <bits/stdc++.h>
 
using namespace std;
 
const int MAXN = 200001;

int n;
vector <int> graph[MAXN];
bool ans;

int dfs(int u, int p) {
    vector <int> child;
    int cnt = 0;

    for (auto v : graph[u]) if (v != p) {
        cnt++;
        auto t = dfs(v, u);
        if (t != 0) {
            child.push_back(t);
        }
    }

    if (child.size() > 1) {
        ans = true;
        return 0;
    }

    int type = cnt == 0 ? 0 : (cnt + 2) % 3;
    vector <int> ord;
    if (type) {
        ord.push_back(type);
    }
    if (!child.empty()) {
        ord.push_back(child[0]);
    }
    if (ord.size() == 2 && ord[0] + ord[1] != 3) {
        ans = true;
        return 0;
    }
    if (ord.size() >= 1) return ord[0];
    return 0;
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int t;
    cin >> t;

    while (t--) {
        cin >> n;
        for (int i = 1; i <= n; ++i) {
            graph[i].clear();
        }
 
        for (int i = 0; i < n - 1; ++i) {
            int u, v;
            cin >> u >> v;
            graph[u].push_back(v);
            graph[v].push_back(u);
        }

        ans = false;
        int res = dfs(1, 1);
        if (res == 2) {
            ans = true;
        }
     
        cout << (ans ? "YES" : "NO") << '\n';
    }

    return 0;
}

Greedy won’t work for this problem, so let’s use dynamic programming instead.

First, we construct the incantations from left to right. Let $$$dp[x][y]$$$ denote the maximum number of additional incantations we can construct if the last incantation is the substring $$$S[x\ldots y]$$$.

For the next incantation, we must find a substring $$$S[y+1\ldots z]$$$ that is lexicographically greater than $$$S[x\ldots y]$$$.

Notice that if $$$S[y+1\ldots z]$$$ is lexicographically greater than $$$S[x\ldots y]$$$, then every longer substring starting at $$$y+1$$$ is also greater. Therefore, we only need to find the smallest possible $$$z$$$.

To find this position efficiently, we calculate the longest common prefix of $$$S[x\ldots y]$$$ and the suffix starting at $$$y+1$$$.

Let $$$lcp[i][j]$$$ denote the length of the longest common prefix of the suffixes $$$S[i\ldots]$$$ and $$$S[j\ldots]$$$. We can calculate it using another dynamic programming:

Using this table, we can find the first valid ending position $$$z$$$ in $$$O(1)$$$ time by comparing the first differing characters. If the entire current incantation matches the beginning of the next suffix, we simply extend the next incantation by at least one character.

After finding (z), we need to find the maximum DP value among all possible ending positions of the next incantation. Checking every position separately would be too slow, so we precompute suffix maximums to answer these queries in (O(1)) time.

Finally, we store the chosen transitions and backtrack through them to reconstruct the spellbook.

The overall time complexity is $$$O(n^2)$$$, where $$$n=|S|$$$, and the space complexity is also $$$O(n^2)$$$.

#include <bits/stdc++.h>
using namespace std;

int dp[5001][5001], lcp[5001][5001];

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

    string s;
    cin >> s;
    int n = s.size();
    s = " " + s;

    for (int i = n; i >= 1; i--) {
        for (int j = n; j >= 1; j--) {
            if (s[i] == s[j])
                lcp[i][j] = 1 + lcp[i + 1][j + 1];
        }
    }

    for (int i = n; i >= 1; i--) {
        for (int j = n; j >= i; j--) {
            int len = min(j - i + 1, lcp[i][j + 1]);

            if (j + 1 + len <= n) {
                if (len == j - i + 1 ||
                    s[i + len] < s[j + 1 + len]) {
                    dp[i][j] = dp[j + 1][j + 1 + len];
                }
            }

            dp[i][j] = max(dp[i][j] + 1, dp[i][j + 1]);
        }
    }

    cout << dp[1][1] << '\n';

    int i = 1, j = 1, remaining = dp[1][1];

    while (i <= n) {
        if (dp[i][j] > remaining) {
            j++;
        } else {
            cout << s.substr(i, j - i + 1) << '\n';
            remaining--;
            i = j + 1;
            j = i;
        }
    }

    return 0;
}