Hi all, I'm planning to do a stream this weekend (edit: now in the past), focusing on problems in Number Theory.

I've put together a mashup of 12 problems (ordered by difficulty), which is available here: https://codeforces.me/contestInvitation/8305134301b2607eaa233fc638ac5fbf91a5c82e (thanks jschnei and neal for problem suggestions and feedback).

Try to solve them on your own! We'll go over those problems & their solutions on stream. Also, feel free to suggest any other number theory problems or concepts you'd be interested in, and I'll try to take a look on stream.

My stream can be found in the sidebar, and you can also follow me directly on Twitch. I upload my streams (afterwards) to YouTube as well, so feel free to subscribe there if you'd like to watch them later!

Edit: I've published this (with timestamps) on YouTube: https://youtu.be/VoD5Mt3Bdec

Hi all, Atcoder Beginner Contest 183 was today. I wrote an unofficial English editorial. Hope it helps!

A: ReLU

We simply write an if statement.

Runtime: $$$\mathcal{O}(1)$$$.

int x; re(x);
if (x < 0) {
    x = 0;
}
ps(x);

B: Billiards

The trick is to reflect $$$G$$$ over the x-axis, so the desired path becomes a straight line. Then it's a matter of simple math to compute the intersection of that line with the x-axis.

Runtime: $$$\mathcal{O}(1)$$$.

double x1, y1, x2, y2;
re(x1, y1, x2, y2);
y2 *= -1;
double r = -y2 / (y1 - y2);
double ans = x1 * r + x2 * (1 - r);
ps(ans);

C: Travel

Because $$$N$$$ is so small, we can simply try all $$$(N-1)!$$$ orderings ($$$7! = 5040$$$), and it takes $$$O(N)$$$ time to evaluate each one. (Why $$$(N-1)!$$$? Well, we know we have to start at $$$1$$$ and end at $$$1$$$, so we can order the $$$N-1$$$ other cities in any order.)

Runtime: $$$\mathcal{O}(N!)$$$.

vector<int> p(n - 1);
iota(all(p), 1);
int ans = 0;
do {
    int cost = 0;
    for (int i = 0; i + 1 < sz(p); i++) {
        int u = p[i];
        int v = p[i + 1];
        cost += t[u][v];
    }
    cost += t[0][p.front()];
    cost += t[p.back()][0];
    ans += (cost == k);
} while (next_permutation(all(p)));

D: Water Heater

We can simulate the process over all periods of time, and check whether there exists a time such that the total water needed at that time is strictly greater than $$$W$$$.

To avoid checking every person at every time step, we can note that each person changes the "current water neded" exactly twice: they add $$$P_i$$$ at time $$$S_i$$$, and they subtract $$$P_i$$$ at time $$$T_i$$$. We can sort these $$$2N$$$ events by time (taking care to put the subtraction events before addition events at the same time, because $$$T_i$$$ is exclusive), then iterate through them.

Runtime: $$$\mathcal{O}(N \log N)$$$.

int n, w; re(n, w);
vector<pair<int, int>> events;
for (int i = 0; i < n; i++) {
    int s, t, p;
    re(s, t, p);
    events.push_back({s, p});
    events.push_back({t, -p});
}
sort(all(events));

ll balance = 0;
for (auto &i : events) {
    balance += i.second;
    if (balance > w) {
        ps(NO);
        return 0;
    }
}
ps(YES);

E: Queen on Grid

We could solve this in $$$O(H \cdot W \cdot \max(H, W))$$$ using a classic DP (for each cell, traverse all cells that could have led to it -- this has $$$HW$$$ states, and the transition takes $$$O(\text{number of previous candidate cells})$$$, which is $$$O(\max(H,W))$$$).

However, because $$$H, W \le 2000$$$, this is too slow. Let's speed up the transition to $$$O(1)$$$, so this will run in time.

Let's just discuss how to optimize the transition for horizontal moves (it's the same for vertical and diagonal). The simplest way to do the transition (but too slow) is to loop over all cells to the left of your current cell, then add their counts to your current cell's count. Instead, we'll simply track the cumulative sum of this, which can be updated in $$$O(1)$$$ and then read in $$$O(1)$$$, instead of $$$O(W)$$$.

Runtime: $$$\mathcal{O}(HW)$$$.

int h, w; re(h, w);
vector<vector<bool>> obs(h, vb(w));
for (int i = 0; i < h; i++) {
    string input; re(input);
    for (int j = 0; j < w; j++)
        obs[i][j] = input[j] == '#';
}

vector<vector<modnum>> dp(h, vector<modnum>(w));
vector<vector<modnum>> horiz(h, vector<modnum>(w));
vector<vector<modnum>> vert(h, vector<modnum>(w));
vector<vector<modnum>> diag(h, vector<modnum>(w));

dp[0][0] = horiz[0][0] = vert[0][0] = diag[0][0] = 1;
obs[0][0] = true; // that way we don't modify (0, 0) any further in the loop
for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
        if (obs[i][j])
            continue;

        if (i - 1 >= 0) {
            dp[i][j] += vert[i - 1][j];
            vert[i][j] = vert[i - 1][j];
        }
        if (j - 1 >= 0) {
            dp[i][j] += horiz[i][j - 1];
            horiz[i][j] = horiz[i][j - 1];
        }
        if (i - 1 >= 0 && j - 1 >= 0) {
            dp[i][j] += diag[i - 1][j - 1];
            diag[i][j] = diag[i - 1][j - 1];
        }

        vert[i][j] += dp[i][j];
        horiz[i][j] += dp[i][j];
        diag[i][j] += dp[i][j];
    }
}

ps(dp.back().back());

F: Confluence

The core idea is that there are two $$$O(NQ)$$$ solutions (too slow!). Let's call type 1 queries "updates" and type 2 queries "queries".

First option: we can use a union-find data structure to maintain the groups. To query, loop over all members of class $$$y$$$ and check if they're in the same group as $$$x$$$. This takes $$$O(1)$$$ to process updates, but $$$O(\text{size of class})$$$ to process queries, which is up to $$$O(N)$$$.

Alternatively: we can use one union-find per class, to maintain the count of students from that class in each group. To query, simply output the already-computed count from that union-find. To update, we have to update all the union-finds, so this takes $$$O(\text{number of classes})$$$, which is up to $$$O(N)$$$. This also takes $$$O(N \cdot (\text{number of classes}))$$$ to initialize.

The key observation here is that the first solution is really good if we have small class sizes, and the second solution is really good if we have a small number of classes. As a result, we'll use a classic trick: use the first solution for small classes, and the second solution for big classes. If we define a "big" class as having at least $$$K$$$ students, the number of big classes is bounded by $$$N/K$$$. This means we can solve the problem in $$$O(N \cdot (N/K) + \max(K, N/K) \cdot Q)$$$. If we set $$$K = \sqrt{N}$$$, then we have a solution in $$$O((N+Q) \sqrt{N} )$$$, which is fast enough if you're careful with the constant factor.

Runtime: $$$\mathcal{O}((N + Q) \sqrt N)$$$.

int n, q; re(n, q);
const int k = 400; // approximately sqrt(2e5)

vector<int> classes(n);
re(classes);
trav (i, classes)
    i--;
vector<vector<int>> students(n);
for (int i = 0; i < n; i++)
    students[classes[i]].push_back(i);
vector<bool> isbig(n);
vector<int> big;
for (int i = 0; i < n; i++) {
    if (sz(students[i]) > k) {
        big.pb(i);
        isbig[i] = true;
    }
}

union_find uf(n);

vector<vector<int>> sizes(n);
for (int i : big) {
    sizes[i] = vector<int>(n);
    for (int j : students[i])
        sizes[i][j] = 1;
}

for (int query = 0; query < q; query++) {
    int type;
    re(type);
    if (type == 1) {
        int a, b;
        re(a, b);
        a--, b--;

        int u = uf.rep(a), v = uf.rep(b);
        bool should = uf.unio(a, b);
        if (should) {
            int r = uf.rep(u);
            int s = u ^ v ^ r;
            for (int i : big) {
                sizes[i][r] += sizes[i][s];
            }
        }
    } else if (type == 2) {
        int x, y;
        re(x, y);
        x--, y--;

        int ans = 0;
        if (isbig[y])
            ans += sizes[y][uf.rep(x)];
        else {
            int r = uf.rep(x);
            for (int j : students[y]) {
                ans += uf.rep(j) == r;
            }
        }
        ps(ans);
    }
}

You can see my submissions here as well, if you prefer that to reading the code in the blog.

Thanks for reading! Let me know if you have any questions or feedback for me.

Hi all, Atcoder Beginner Contest 174 was today. I wrote an unofficial English editorial. Hope it helps!

A: Air Conditioner

We simply write an if statement.

Runtime: $$$\mathcal{O}(1)$$$.

int x = in.nextInt();
out.println(x >= 30 ? "Yes" : "No");

B: Distance

We can loop through all the points, and test each one. To avoid potential rounding errors, it's simplest to check $$$x^2 + y^2 \le D^2$$$, so we can keep everything in integers.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
long d = in.nextInt();
d = d * d;
Point[] points = Point.readPoints(in, n);

int answer = 0;
for (Point p : points) {
    if (p.x * p.x + p.y * p.y <= d)
        answer++;
}

out.println(answer);

C: Repsept

First, we simplify: if $$$k$$$ is a multiple of $$$7$$$, then we can look for a number like $$$1111$$$ that's a multiple of $$$k/7$$$. Otherwise, using $$$7777$$$ instead of $$$1111$$$ doesn't help us.

If you consider the sequence of numbers to check as $$$a_i = 10 a_{i-1} + 1$$$ modulo $$$k$$$, there is guaranteed to be a solution within $$$k$$$ steps if and only if $$$\mathrm{gcd}(10, k) = 1$$$. So we can check for impossibility using this fact, and then iterate through the choices otherwise and check them all until we find the smallest one that works.

Take care to test locally on $$$k=1$$$ specifically, it's easy to get this wrong with a loop.

Runtime: $$$\mathcal{O}(k)$$$.

int k = in.nextInt();
if (k % 7 == 0)
    k /= 7;

if (gcd(10, k) != 1) {
    out.println(-1);
    return;
}

long v = 1 % k;
int len = 1;
while (v != 0) {
    v = (10 * v + 1) % k;
    len++;
}

out.println(len);

D: Alter Altar

A few quick observations:

• At the end, we want a sequence that looks like RRRRWWWWWWW (some number of red stones, followed by some number of white stones, and one of these could possibly be zero).
• If we want to change a red stone's color and a white stone's color, this takes 1 step as we can swap them.

Now, we can simply try all values of the final number of red stones from $$$0$$$ to $$$N$$$ (let's call this number $$$R$$$). For a given value of $$$R$$$, the cost is given as

If we compute prefix sums to help us compute this, we can test one value in $$$O(1)$$$, so testing them all will run in time.

Bonus: this function is convex, so we can minimize it using ternary search.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
char[] c = in.next().toCharArray();

int[] w = new int[n + 1];
for (int i = 0; i < n; i++) {
    w[i + 1] = w[i] + (c[i] == 'W' ? 1 : 0);
}

IntUnaryOperator moves = left -> {
    int wrongLeft = w[left];
    int wrongRight = (n - left) - (w[n] - w[left]);

    return wrongLeft + wrongRight - Math.min(wrongLeft, wrongRight);
};

int answer = ternarySearch(moves, 0, n);
out.println(moves.applyAsInt(answer));

E: Logs

It's hard to figure out which logs to cut greedily, but given a value $$$L$$$ for the final answer, we can easily check whether it's attainable in $$$O(N)$$$.

In order to do this, we loop through all the logs and cut off pieces of exactly length $$$L$$$ from them, until they are length $$$L$$$ or shorter. So for a log of length $$$x$$$, this takes $$$\lfloor(x-1)/L \rfloor$$$ steps.

It's also easy to see intuitively that the cost is non-increasing as $$$L$$$ increases, so we can binary search to find the smallest length $$$L$$$ so that the numbers of steps is $$$\le k$$$.

Runtime: $$$\mathcal{O}(N \log L)$$$, where $$$L$$$ is the maximum possible length of a log.

int n = in.nextInt();
int k = in.nextInt();
int[] a = in.readIntArray(n);

IntPredicate cuts = l -> {
    if (l <= 0)
        return false;
    long c = 0;
    for (int x : a) {
        c += (x - 1) / l;
    }
    return c <= k;
};

out.println(binarySearch(cuts));

F: Range Set Query

This is a standard algorithm, which I didn't figure out myself but learned some time ago by web search, from the following link: https://www.geeksforgeeks.org/queries-number-distinct-elements-subarray/

To implement it efficiently without bugs, I copy/pasted from my submission for a previous Codeforces problem (85732941).

Runtime: $$$\mathcal{O}((N + Q) \log N)$$$.

Time to write: $$$O(1)$$$.

int n = in.nextInt(), queryCount = in.nextInt();
int[] c = in.readIntArray(n);
for (int i = 0; i < n; i++) {
    c[i]--;
}
Pii[] queries = new Pii[queryCount];
for (int i = 0; i < queryCount; i++) {
    queries[i] = Pii.of(in.nextInt() - 1, in.nextInt());
}

List<Integer>[] byEnd = Util.arrayOfLists(n + 1);
for (int i = 0; i < queryCount; i++) {
    byEnd[queries[i].second].add(i);
}

IntSumSegmentTree st = new IntSumSegmentTree(n);
int[] last = new int[n];
Arrays.fill(last, -1);

int[] answer = new int[queryCount];
for (int i = 0; i < n; i++) {
    if (last[c[i]] != -1) {
        st.update(last[c[i]], 0);
    }

    st.update(i, 1);
    last[c[i]] = i;

    for (int j : byEnd[i + 1]) {
        answer[j] = st.query(queries[j].first, queries[j].second);
    }
}

for (int x : answer)
    out.println(x);

You can see my submissions here as well, if you prefer that to reading the code in the blog.

Thanks for reading! Let me know if you have any questions or feedback for me.

Hi all, Atcoder Beginner Contest 171 was today. I wrote an unofficial English editorial. Hope it helps!

A: αlphabet

We simply write an if statement.

Runtime: $$$\mathcal{O}(1)$$$.

char x = in.next().charAt(0);

if (x >= 'A' && x <= 'Z')
    out.println("A");
else
    out.println("a");

B: Mix Juice

We want to buy the cheapest fruits, so we can sort the array and then pick the first $$$K$$$ fruits.

Runtime: $$$\mathcal{O}(N \log N)$$$.

int n = in.nextInt(), k = in.nextInt();
int[] p = in.readIntArray(n);

Util.safeSort(p); // this shuffles and sorts to avoid hacks

long answer = 0;
for (int i = 0; i < k; i++) {
   answer += p[i];
}

out.println(answer);

C: One Quadrillion and One Dalmatians

The trick is to look at the letters starting from the end of the names. You can see the last letter repeats every $$$26$$$ names. The second-to-last letter repeats every $$$26^2$$$ names (after you skip the first $$$26$$$ names that are too short). The third-to-last letter repeats every $$$26^3$$$ names (after you skip the first $$$26 + 26^2$$$ names that are too short).

This gives rise to a solution where we extract the last letter first by setting $$$N = N-1$$$, then looking at the value modulo $$$26$$$. Then to shift everything over, we divide by $$$26$$$, and repeat the whole process as necessary. The subtraction of $$$1$$$ at the beginning of the loop accounts for skipping all the necessary shorter strings.

Take care to reverse the string.

Runtime: $$$\mathcal{O}(\log_{26} N)$$$, or equivalently $$$\mathcal{O}(L)$$$ where $$$L$$$ is the length of the answer.

long n = in.nextLong();

StringBuilder answer = new StringBuilder();
while (n > 0) {
    n--;
    answer.append((char) ('a' + n % 26));
    n /= 26;
}

out.println(answer.reverse());

D: Replacing

A few quick observations:

• The order of $$$A$$$ doesn't matter.
• Once two elements of $$$A$$$ are equal, they will never diverge.

This motivates us to think about the counts of each element in $$$A$$$, rather than the elements themselves. Because the elements are small (less than $$$10^5$$$), we can easily store the counts in an array.

In order to process a query, we move everything from $$$B_i$$$ to $$$C_i$$$ and update our sum appropriately. This means adding $$$\mathrm{count}[B_i]$$$ to $$$\mathrm{count}[C_i]$$$ and setting $$$\mathrm{count}[B_i]$$$ to zero, as well as changing our running sum by $$$\mathrm{count}[B_i] (C_i - B_i)$$$.

Take care to use $$$\mathrm{count}[B_i]$$$ before you set it to zero.

Runtime: $$$\mathcal{O}(N + Q + M)$$$, where $$$M = 10^5$$$, the largest possible element value.

int n = in.nextInt();
int[] a = in.readIntArray(n);
long sum = 0;

long[] count = new long[MAX];
for (int i : a) {
    sum += i;
    count[i]++;
}

int q = in.nextInt();
for (int query = 0; query < q; query++) {
    int b = in.nextInt(), c = in.nextInt();

    sum -= count[b] * b;
    sum += count[b] * c;
    count[c] += count[b];
    count[b] = 0;

    out.println(sum);
}

E: Red Scarf

The key observation (and super useful fact about xor in general) is that $$$x \oplus x = 0$$$ for any $$$x$$$. This means a lot cancels out.

Let $$$c_i$$$ be a valid original array of cats (so, this is our answer we want to output).

Let's first rewrite $$$a_i$$$ using our observation above. Let $$$X$$$ be the xor of all elements in $$$c$$$. Then $$$a_i = X \oplus c_i$$$.

So, we can simply compute $$$c_i = X \oplus a_i$$$ and we are done.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
int[] a = in.readIntArray(n);

int xor = 0;
for (int x : a)
    xor ^= x;

int[] answer = new int[n];
for (int i = 0; i < n; i++) {
    answer[i] = xor ^ a[i];
}

out.println(Util.join(answer));

F: Strivore

Let's consider the end result of this process and what it can look like. It will be a string (let's call it $$$T$$$) of length $$$N = K + |S|$$$ that contains $$$S$$$ as a subsequence.

We'll use $$$Z = 26$$$ for readability.

In order to count these, let's first note that there are $$$Z^N$$$ total strings of length $$$N$$$, and we want to remove all the strings that don't contain $$$S$$$ as a subsequence (this is a bit easier than counting the strings that do contain $$$S$$$, because a string could contain $$$S$$$ multiple times and be overcounted).

Let's consider the length of the maximum prefix of $$$S$$$ that's a subsequence of our string $$$T$$$. Going from left to right within $$$T$$$, we can check when this increases -- it increases from $$$X$$$ to $$$X+1$$$ when our current character is $$$S_{X+1}$$$. In order for it to not increase, we need our current character to not be $$$S_{X+1}$$$.

So let's say this length is exactly $$$M$$$. Then it must have increased from $$$0$$$ to $$$M$$$ (one at a time, of course), so we pick $$$M$$$ indices in $$$T$$$ to be the indices where it increases, and those have only one choice of character. The remaining characters have $$$Z-1$$$ choices. So for length $$$M$$$, we have $$$\binom{N}{M} (Z-1)^{N-M}$$$ possibilities.

In order to not contain $$$S$$$ as a subsequence, $$$M$$$ can be anything from $$$0$$$ to $$$|S| - 1$$$. Now, recall these are the strings we don't want, so we subtract them from the total.

Therefore, our final answer is $$$\displaystyle Z^N - \left[\sum_{M=0}^{|S|-1} \binom{N}{M} (Z-1)^{N-M}\right]$$$.

Runtime: $$$\mathcal{O}(K + |S|)$$$.

final int Z = 26;
int k = in.nextInt();
char[] s = in.next().toCharArray();
int l = s.length;
int n = k + l;

// this is all just some precomputation, so we can
// get powers and nCr in O(1) later
mod.fact(n);
long[] powers = new long[n + 1];
powers[0] = 1;
for (int i = 0; i + 1 < powers.length; i++) {
    powers[i + 1] = mod.mult(powers[i], Z - 1);
}

long answer = mod.pow(Z, n);
for (int match = 0; match < l; match++) {
    long val = mod.mult(powers[n - match], mod.ncr(n, match));
    answer = mod.subtract(answer, val);
}

out.println(answer);

You can see my submissions here as well, if you prefer that to reading the code in the blog.

Thanks for reading! Let me know if you have any questions or feedback for me.

Hi all, Atcoder Beginner Contest 146 was today. I wrote an unofficial English editorial. Hope it helps!

A: Can't Wait for Holiday

We simply find the index in the week, and print $$$7-i$$$ (taking care to order the days correctly so Sunday gives us $$$7$$$ as the output).

Runtime: $$$\mathcal{O}(1)$$$.

String w = "SUNMONTUEWEDTHUFRISAT"; // this is sort of a bad way to do this
String s = in.next();
int i = w.indexOf(s) / 3;
out.println(7 - i);

B: ROT N

We can simply loop through the string and increment each character by $$$N$$$, taking care to subtract $$$26$$$ if we go past the end of the alphabet.

Runtime: $$$\mathcal{O}(|S|)$$$.

int n = in.nextInt();
char[] s = in.next().toCharArray();
for (int i = 0; i < s.length; i++) {
    s[i] += n;
    if (s[i] > 'Z')
        s[i] -= 26;
}
out.println(String.valueOf(s));

C: Buy an Integer

For a given length $$$L$$$, we can consider the maximum integer we can afford of that length (this is easy to calculate because given a fixed $$$d(N)$$$, the cost is monotonic in $$$N$$$). There are only a small number of possible lengths, so we may loop through them.

Note that if any number greater than $$$10^9$$$ is affordable, then $$$10^9$$$ is as well (since it has smaller-or-equal numeric value and length than any number greater than it).

Runtime: $$$\mathcal{O}(\log_{10} X)$$$.

long a = in.nextInt(), b = in.nextInt();
long x = in.nextLong();

long answer = 0;
for (int len = 1; len <= 18; len++) {
    long aBudget = x - b * len;
    long n = aBudget / a;
    if (String.valueOf(n).length() > len)
        continue;
    answer = Math.max(answer, n);
}
answer = Math.min(answer, 1000000000);

out.println(answer);

D: Coloring Edges on Tree

We can construct an optimal coloring by greedily assigning colors to edges while traversing the tree. Given a vertex $$$v$$$, we can assign colors $$$1$$$ through $$$\mathrm{deg}(v)$$$ to its unassigned incident edges (if one of them is already assigned, we leave it as-is). The final answer for total colors is $$$\mathrm{max}_v(\mathrm{deg}(v))$$$.

I implemented this using a BFS through the tree. When visiting a node, at most one of its edges is already assigned (the one we just came from, if we aren't at the root), so we simply skip this one and note that we can't use its color for another edge on this node.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
List<Integer>[] adj = new List[n];
for (int i = 0; i < n; i++)
    adj[i] = new ArrayList<>();
HashMap<Pair<Integer, Integer>, Integer> edgeIndex = new HashMap<>();
for (int i = 0; i < n - 1; i++) {
    int a = in.nextInt() - 1, b = in.nextInt() - 1;
    adj[a].add(b);
    adj[b].add(a);
    edgeIndex.put(Pair.of(a, b), i);
    edgeIndex.put(Pair.of(b, a), i);
}

int[] color = new int[n - 1];
Queue<Pair<Integer, Integer>> q = new ArrayDeque<>();
q.add(Pair.of(0, -1));
while (!q.isEmpty()) {
    Pair<Integer, Integer> p = q.poll();
    int cur = p.first;
    int parent = p.second;
    int colorOfParent = parent == -1 ? 0 : color[edgeIndex.get(Pair.of(cur, parent))];
    int skippedColor = -1, skippedEdge = -1;
    for (int i = 0; i < adj[cur].size(); i++) {
        if (adj[cur].get(i) == parent) {
            skippedColor = i + 1;
            continue;
        }
        if (i + 1 == colorOfParent) {
            skippedEdge = edgeIndex.get(Pair.of(cur, adj[cur].get(i)));
        } else {
            color[edgeIndex.get(Pair.of(cur, adj[cur].get(i)))] = i + 1;
        }
        q.add(Pair.of(adj[cur].get(i), cur));
    }
    if (skippedEdge >= 0)
        color[skippedEdge] = skippedColor;
}
out.println(Util.max(color));
for (int x : color)
    out.println(x);

E: Rem of Sum is Num

First, we can consider all sums modulo $$$K$$$ throughout this problem. Let's weaken the constraints of what we're looking for and consider the number of elements modulo $$$K$$$ for a moment as well.

Notice that if we let $$$S_n = \sum_{i=1}^n A_i$$$ (and $$$S_0 = 0$$$), then the sum of a range $$$[i, j]$$$ (1-indexed) is $$$S_j - S_{i-1}$$$ and the number of elements in it is $$$j - (i-1)$$$.

This leads us to compute $$$P_i = S_i - i$$$ for $$$0 \le i \le n$$$, and we look for pairs $$$i < j$$$ such that $$$P_i \equiv P_j \pmod K$$$ (corresponding to taking the range $$$A_{i+1}, \ldots, A_j$$$).

Note that we actually have one more constraint: we shouldn't actually take the number of elements modulo $$$K$$$, so we need the number of elements modulo $$$K$$$ to be equal to the actual number of elements, so we require $$$j-(i-1) < K$$$.

We can count the pairs by maintaining a hashmap of the counts of values of $$$P_n$$$ within a sliding window of the past $$$K$$$ items (these are our candidates for the first coordinate of our range), and adding the count that match each possible second coordinate of our range. Each update takes constant time, so overall this solution takes linear time.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
int k = in.nextInt();
int[] a = in.readIntArray(n);
for (int i = 0; i < n; i++)
    a[i] %= k;
long[] prefix = new long[n + 1];
for (int i = 0; i < n; i++) {
    prefix[i + 1] = prefix[i] + a[i];
}
for (int i = 0; i <= n; i++) {
    prefix[i] = ((prefix[i] - i) % k + k) % k;
}

HashMap<Long, Integer> count = new HashMap<>();
count.put(0L, 1);
long answer = 0;
for (int i = 1; i <= n; i++) {
    if (i - k >= 0) {
        count.merge(prefix[i - k], -1, Integer::sum);
    }
    answer += count.getOrDefault(prefix[i], 0);
    count.merge(prefix[i], 1, Integer::sum);
}

out.println(answer);

F: Sugoroku

Let's make a few observations.

Observation 1: in order to minimize the number of steps, we should take steps that are as large as possible. We don't have to worry about a smaller step ever being better, because if in one step we can choose to go to either $$$i$$$ or $$$j$$$ (with $$$i < j$$$), then anything reachable from $$$i$$$ is also reachable from $$$j$$$ in the same-or-fewer steps. To prove this, consider an index $$$k > j$$$ that is reachable from $$$i$$$. At some point when traveling from $$$i$$$ to $$$k$$$, we will take a step that passes $$$j$$$. At this point, the endpoint of that step is also reachable from $$$j$$$ in one step (because it's at most $$$M$$$ away from a square that's before $$$j$$$), so we could have gotten here at least as fast from $$$j$$$.

Observation 2: in order to find the lexicographically-smallest sequence for a given number of steps, we need to put the largest steps at the end of the sequence. Therefore, we can start at $$$N$$$ and takes steps as large as possible towards $$$0$$$, then reverse the order of the steps at the end for printing.

Observation 3: Let's say we took a (backwards) step from $$$L$$$ to $$$C$$$ just now ($$$L > C$$$), and $$$C$$$ is the minimum index reachable in one step from $$$L$$$ (since we take the biggest steps possible). Then when evaluating the candidates for our next move, we don't need to evaluate anything that's $$$\ge (L-M)$$$, because those are all either unreachable or greater than $$$C$$$. This optimizes our runtime from quadratic to linear, since we only end up looking at each candidate once ever.

We can put this all together to code up a fairly straightforward linear-time solution.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
int m = in.nextInt();
char[] inputS = in.next().toCharArray(); // len n+1
boolean[] s = new boolean[n + 1];
for (int i = 0; i < n; i++) {
    s[i] = inputS[i] == '1';
}

Deque<Integer> answer = new ArrayDeque<>();

int cur = n;
int last = Integer.MAX_VALUE;
outer:
while (cur > 0) {
    for (int i = Math.max(0, cur - m); i < cur && i < last - m; i++) {
        if (!s[i]) {
            answer.addFirst(cur - i);
            last = cur;
            cur = i;
            continue outer;
        }
    }
    out.println(-1);
    return;
}

out.println(Util.join(answer));

Thanks for reading! Let me know if you have any questions or feedback for me.

Hi all, Atcoder Beginner Contest 143 was today. I wrote an unofficial English editorial. Hope it helps!

A: Curtain

The curtains can cover a maximum of $$$2B$$$. So either they don't cover the whole window and the remainder is $$$A-2B$$$, or they do and the remainder is $$$0$$$. So the answer is $$$\max(0, A-2B)$$$.

Runtime: $$$\mathcal{O}(1)$$$.

int a = in.nextInt(), b = in.nextInt();
out.println(Math.max(0, a - 2 * b));

B: TAKOYAKI FESTIVAL 2019

We can simply loop through every pair (taking care not to double count) and add the health points restored.

Runtime: $$$\mathcal{O}(N^2)$$$.

int n = in.nextInt();
int[] d = in.readIntArray(n);

long answer = 0;
for (int i = 0; i < n; i++) {
    for (int j = i + 1; j < n; j++) {
        answer += d[i] * d[j];
    }
}

out.println(answer);

Alternate solution (and code) by Tejs:

The product can be rewritten as $$$\frac{1}{2} \left[ \left(\sum_i d_i\right)^2 - \sum_i d_i^2\right]$$$.

Runtime: $$$\mathcal{O}(N)$$$.

int n;
cin >> n;
vector<int> ds(n);
for (auto& d : ds) {
    cin >> d;
}
int sum_a = 0, sum_b = 0;
for (auto d : ds) {
    sum_a += d;
    sum_b += d*d;
}
cout << (sum_a*sum_a - sum_b)/2 << endl;

C: Slimes

We have to count the number of "runs" of consecutive slimes of the same color. We can do this by looping through the array and checking whether the current slime continues the run, then incrementing our answer when we reach the end of a run.

Runtime: $$$\mathcal{O}(N)$$$.

int n = in.nextInt();
char[] s = in.next().toCharArray();

int answer = 0;
for (int i = 0, j = 0; i < n; i = j) {
    while (j < n && s[i] == s[j])
        j++;

    answer++;
}
out.println(answer);

D: Triangles

The total number of (unordered) triples of sticks is $$$\binom{n}{3} = n(n-1)(n-2)/6$$$. I think it's simplest to count the number of triples that cannot form a triangle, and subtract them. Let's first sort the input array, so if we have a triple $$$(i,j,k)$$$ with $$$i < j < k$$$, we also know $$$L_i \le L_j \le L_k$$$.

Then, we can loop through all pairs $$$(i,j)$$$ and determine how many values of $$$k$$$ make an invalid triple (but still maintaining $$$i<j<k$$$). Note that because we picked the two smaller sticks of our triple, we know the only possible disqualifying condition is $$$L_k \ge L_i + L_j$$$. We can find the smallest such value of $$$k$$$ using binary search, and then subtract $$$n-k$$$ from our answer.

Runtime: $$$\mathcal{O}(N^2 \log N)$$$.

int n = in.nextInt();
Integer[] L = in.readIntegerArray(n);

Arrays.sort(L);

long answer = ((long) n) * (n - 1) * (n - 2) / 6;
for (int i = 0; i < n; i++) {
    for (int j = i + 1; j < n; j++) {
        int sum = L[i] + L[j];

        // binarySearch(condition) finds the first index such that condition is true
        int k = binarySearch(index -> {
            if (index >= n)
                return true;
            if (index < 0)
                return false;
            return L[index] >= sum;
        });

        answer -= n - k;
    }
}

out.println(answer);

E: Travel by Car

First, observe that any path that has $$$K$$$ refills can be split into $$$K+1$$$ paths, each of which can be completed (starting with a full tank of gas) without refilling.

So, we can begin by finding all pairs of cities that can be traveled between on a full tank of gas without refilling. To do this, we can find the shortest path between all pairs of cities (for example, using Floyd-Warshall). Then, we build a new adjacency matrix for these paths, where two cities with distance $$$\le L$$$ have an edge of cost $$$1$$$, and two cities with distance $$$> L$$$ have no edge (or an edge with cost $$$\infty$$$). Then we compute the shortest path between all pairs using this new adjacency matrix, which tells us the minimum number of segments a valid path in the original graph can be broken into.

Then, given a query $$$(s,t)$$$, we simply find the distance in the new shortest-paths matrix and subtract $$$1$$$, and we have our answer (or print $$$-1$$$ if it's unreachable). We have to remember to subtract $$$1$$$ because it takes $$$K$$$ refills for a path of $$$K+1$$$ segments (or equivalently, the first tank of gas is free).

Runtime: $$$\mathcal{O}(N^3)$$$.

int n = in.nextInt();
int m = in.nextInt();
long l = in.nextLong();

long[][] adj = new long[n][n];
for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
        adj[i][j] = i == j ? 0 : INF;
    }
}
for (int i = 0; i < m; i++) {
    int a = in.nextInt() - 1;
    int b = in.nextInt() - 1;
    long c = in.nextLong();

    adj[a][b] = adj[b][a] = c;
}

// Floyd-Warshall
for (int k = 0; k < n; k++) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            adj[i][j] = Math.min(adj[i][j], adj[i][k] + adj[k][j]);
        }
    }
}

long[][] oldAdj = adj.clone();

for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
        adj[i][j] = oldAdj[i][j] <= l ? 1 : INF;
        if (i == j) {
            adj[i][j] = 0;
        }
    }
}

// Floyd-Warshall
for (int k = 0; k < n; k++) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            adj[i][j] = Math.min(adj[i][j], adj[i][k] + adj[k][j]);
        }
    }
}

int q = in.nextInt();
for (int query = 0; query < q; query++) {
    int s = in.nextInt() - 1;
    int t = in.nextInt() - 1;

    long answer = adj[s][t];
    if (answer >= INF) {
        answer = -1;
    } else {
        answer--; // first fill free
    }
    out.println(answer);
}

F: Distinct Numbers

The key observation is that if you pick $$$M$$$ distinct numbers and remove all instances of them, and you have $$$S$$$ total numbers left, the maximum number of operations you can do is bounded above by $$$S/(K-M)$$$. Then, note that picking the $$$M$$$ most frequent distinct numbers gives the tightest bound (so instead of all subsets, we need to consider a lot less stuff).

So, we can count the instances of each number, so we end up with an array $$$C$$$ containing these frequencies. We can sort these frequencies, then for each value of $$$M$$$, compute the number of remaining numbers after we remove all instances of the top $$$M$$$ (let's call this $$$S_M$$$).

The function $$$f_K(M) = \frac{S_M}{K-M}$$$ is first decreasing, then increasing, so we can find its minimum by binary searching to find the first $$$M$$$ for which it increases. This gives us the tightest bound available. Once we do that, we simply need to check against the simple bound $$$N/K$$$ and return our answer.

Runtime: $$$\mathcal{O}(N \log {N})$$$.

int n = in.nextInt();
int[] a = in.readIntArray(n);
Integer[] count = new Integer[n];
Arrays.fill(count, 0);

for (int x : a) {
    count[x - 1]++;
}

Arrays.sort(count, Comparator.reverseOrder());

int[] remain = new int[n];
remain[0] = n - count[0];
for (int i = 1; i < n; i++) {
    remain[i] = remain[i - 1] - count[i];
}

out.println(n); // k=1
for (int kLoop = 2; kLoop <= n; kLoop++) {
    int k = kLoop;
    int bound = n / k;
    
    // finds the first i s.t. the condition is true
    int index = BinarySearch.binarySearch((i) -> {
        if (k - 1 - i <= 0)
            return true;
        if (i <= 0)
            return false;

        double boundCur = ((double) remain[i]) / (k - 1 - i);
        double boundPrev = ((double) remain[i - 1]) / (k - 1 - (i - 1));
        return boundPrev < boundCur;
    });

    index--;
    bound = Math.min(bound, remain[index] / (k - 1 - index));
    out.println(bound);
}

You can see my submissions here as well, if you prefer that to reading the code in the blog.

Thanks for reading! Let me know if you have any questions or feedback for me.

Hi all, Atcoder Beginner Contest 132 was today. Atcoder only publishes Japanese editorials for beginner contests, so I wrote an unofficial English editorial. Hope it helps!

A: Fifty-Fifty

After sorting the input, a "good" string will look like "AABB". Therefore, we simply sort and check for this format.

char[] s = in.next().toCharArray();
Arrays.sort(s);
out.println(s[0] == s[1] && s[2] == s[3] && s[1] != s[2] ? "Yes" : "No");

B: Ordinary Number

It suffices to simply check every range of 3 and see if its middle element should be counted. A simple way is to sort each range of 3 and check that the middle element remains in the middle.

Runtime: $$$\mathcal{O}(n)$$$.

int n = in.nextInt();
int[] p = new int[n];
for (int i = 0; i < n; i++)
    p[i] = in.nextInt();

int answer = 0;
for (int i = 1; i + 1 < n; i++) {
    int[] x = {p[i - 1], p[i], p[i + 1]};
    Arrays.sort(x);
    if (p[i] == x[1])
        answer++;
}
out.println(answer);

C: Divide the Problems

If we first sort the problems by difficulty, the necessary condition is equivalent to "the first half of the problems will be for ABCs, and the second half will be for ARCs". If we assign $$$B$$$ to the hardest ABC problem (problem $$$(N/2)$$$, after sorting) and $$$R$$$ to the easiest ARC problem (problem $$$(N/2+1)$$$, after sorting), we need to select $$$K$$$ such that $$$B < K \le R$$$.

Therefore, we have $$$R-B$$$ choices for K. (Note: if $$$R$$$ and $$$B$$$ are equal, there is no choice of $$$K$$$ that splits the problems in half.)

Runtime: $$$\mathcal{O}(N \log N)$$$.

int n = in.nextInt();
Integer[] d = new Integer[n];
for (int i = 0; i < n; i++)
    d[i] = in.nextInt();

Arrays.sort(d);
out.println(d[n / 2] - d[n / 2 - 1]);

D: Blue and Red Balls

In order for Takahashi to take exactly $$$i$$$ moves, we need $$$i$$$ separate "buckets" of blue balls. Let us say we have $$$R = N-K$$$ red balls and $$$B = K$$$ blue balls. Then, we may separately compute $$$X$$$ as the number of ways to order the $$$R$$$ red balls and the $$$i$$$ buckets of blue balls, and compute $$$Y$$$ as the number of ways to distribute the $$$B$$$ blue balls into the $$$i$$$ buckets. The output for $$$i$$$ moves is then $$$X \cdot Y$$$.

Let us compute $$$X$$$. If we imagine the $$$R$$$ red balls already in a row, there are $$$R+1$$$ spaces to put a bucket of blue balls (it can go left of all the balls, right of all the balls, or in the $$$R-1$$$ gaps between consecutive red balls). Therefore, $$$X = \binom{R+1}{i}$$$.

Let us now compute $$$Y$$$. To distribute $$$B$$$ blue balls into $$$i$$$ buckets, we can imagine all the blue balls in a row and insert $$$i-1$$$ dividers into the $$$B-1$$$ spaces between the balls to separate them into buckets. Therefore, $$$Y = \binom{B-1}{i-1}$$$.

Note: computing binomial coefficients can be done for this problem either by precomputing Pascal's Triangle in $$$O(N^2)$$$ time or by precomputing factorials and their inverses, since $$$N \le 2000$$$.

int n = in.nextInt(), k = in.nextInt();

int r = n - k, b = k;
Mod107 m = new Mod107();

for (int i = 1; i <= k; i++) {
   ModularNumber<Mod107> x = m.ncr(r + 1, i);
   ModularNumber<Mod107> y = m.ncr(b - 1, i - 1);
   out.println(x.mult(y));
}

E: Hopscotch Addict

Repeating ken-ken-pa $$$k$$$ times corresponds to following a path of length exactly $$$3k$$$ in the graph. Therefore, we need to find the shortest path from $$$S$$$ to $$$T$$$ whose length is a multiple of 3.

In order to solve this, we create a new graph $$$G'$$$ with $$$3N$$$ vertices and $$$3M$$$ edges. For each vertex $$$v_i$$$ in the original graph $$$G$$$, we create three vertices $$$v_0$$$, $$$v_1$$$, and $$$v_2$$$ in $$$G'$$$. For every edge $$$(u, v)$$$ in $$$G$$$, we create edges $$$(u_0, v_1)$$$, $$$(u_1, v_2)$$$, and $$$(u_2, v_0)$$$ in $$$G'$$$.

In this way, a path from $$$u_i$$$ to $$$v_j$$$ in $$$G'$$$ corresponds to a path in $$$G$$$ whose length is congruent to $$$j-i \text{ (mod 3)}$$$.

Therefore, we can run a BFS on $$$G'$$$ to find the shortest path from $$$S_0$$$ to $$$T_0$$$, which corresponds to the shortest path with length a multiple of 3 in $$$G$$$. If we find such a path, we simply divide its length by 3 to get the number of ken-ken-pa.

Runtime: $$$\mathcal{O}(N + M)$$$.

int n = in.nextInt();
int m = in.nextInt();

List<Integer>[] adj = new List[3 * n];
for (int i = 0; i < 3 * n; i++) {
    adj[i] = new ArrayList<>();
}

for (int i = 0; i < m; i++) {
    int u = in.nextInt() - 1;
    int v = in.nextInt() - 1;

    // Create G'
    adj[u].add(v + n);
    adj[u + n].add(v + 2 * n);
    adj[u + 2 * n].add(v);
}

// BFS from s to t
int s = in.nextInt() - 1, t = in.nextInt() - 1;
Queue<Integer> q = new ArrayDeque<>();
q.add(s);
BitSet visited = new BitSet();
visited.set(s);
int[] dist = new int[3 * n];
dist[s] = 0;
while (!q.isEmpty()) {
    int cur = q.poll();
    for (int next : adj[cur]) {
        if (visited.get(next))
            continue;

        visited.set(next);
        dist[next] = dist[cur] + 1;
        q.add(next);
    }
}

// Divide path length by 3 to get number of ken-ken-pa.
out.println(visited.get(t) ? dist[t] / 3 : -1);

F: Small Products

Let's divide the numbers from $$$1$$$ to $$$N$$$ into "small" and "big" numbers. We call "small" numbers those that are at most $$$\sqrt{N}$$$, and "big" numbers those that are greater than $$$\sqrt{N}$$$.

This means that any two small numbers can be adjacent, and no two big numbers can be adjacent. When putting a small number $$$s$$$ and a big number $$$b$$$ adjacent, we require $$$s \cdot b \le N$$$.

Thus, we can split the possible values of $$$b$$$ into $$$\sqrt{N}$$$ buckets $$$B_i$$$ based on the maximal small number they can be paired with (for example, if $$$N=10$$$, the big numbers $$$4$$$ and $$$5$$$ can be paired with the small number $$$2$$$, but not with $$$3$$$, so they would go in $$$B_2$$$).

When we have built a partial sequence, we only care about the last value in the sequence when deciding how to build the rest of the sequence. Moreover, we actually only care about the $$$2\sqrt{N}$$$ categories the last element can fall into (either a small number or a bucket of big numbers).

Let $$$s(i, j)$$$ be the number of sequences of length $$$i$$$ ending with a small number $$$j$$$. Let $$$b(i, j)$$$ be the number of sequences of length $$$i$$$ ending with a big number in $$$B_i$$$.

A small number $$$j$$$ can be placed after any other small number, or after a big number in $$$B_k$$$ where $$$k \ge i$$$.

$$$\displaystyle s(i,j) = \sum_{k=1}^{\sqrt{N}} s(i-1,k) + \sum_{k=j}^{\sqrt{N}} b(i-1,k)$$$

A big number $$$j$$$ can be placed after any small number $$$k \le j$$$. We have $$$|B_j| = \frac{N}{j} - \frac{N}{j+1}$$$ possibilities for this big number.

$$$\displaystyle b(i,j) = \left( \frac{N}{j} - \frac{N}{j+1} \right) \sum_{k=1}^{j} s(i-1,k)$$$

This directly leads to an $$$\mathcal{O}(k \sqrt{N}^2)$$$ dynamic programming solution, which can be sped up to $$$\mathcal{O}(k \sqrt{N})$$$ by precomputing $$$\sum_{k=1}^j s(i-1,j)$$$ and $$$\sum_{k=j}^{\sqrt{N}} b(i-1,k)$$$ for all j, after we compute everything for $$$i-1$$$ and before we compute everything for $$$i$$$.

Runtime: $$$\mathcal{O}(k \sqrt{N})$$$.

long n = in.nextLong();
int k = in.nextInt();

// sqrtN is the smallest number that when squared is greater than n
int sqrtN = (int) BinarySearch.binarySearchLong(i -> i * i > n);

// dpSmall[i][j] is the number of sequences of length i where the last number is j.
long[][] dpSmall = new long[k + 1][sqrtN];
// dpBig[i][j] is the number of sequences of length i where the last number is in bucket B_j.
long[][] dpBig = new long[k + 1][sqrtN];

// The first number can be anything, which we represent by having exactly 1 seq of length 0 that
// "ends" with 1.
dpSmall[0][1] = 1;

for (int i = 0; i < k; i++) {
    for (int j = 1; j < sqrtN; j++) {
        dpSmall[i][j] += dpSmall[i][j - 1];
        dpSmall[i][j] %= MOD;
    }
    for (int j = sqrtN - 2; j > 0; j--) {
        dpBig[i][j] += dpBig[i][j + 1];
        dpBig[i][j] %= MOD;
    }

    // now dpSmall[i][...] is replaced by partial sums
    // now dpBig[i][...] is replaced by partial sums

    for (int j = 1; j < sqrtN; j++) {
        // If the previous number is anything <= j, we have (n/j - n/(j+1)) choices for this number.
        dpBig[i + 1][j] = dpSmall[i][j] * (n / j - n / (j + 1));

        // We don't want to count a number in dpBig that could be paired with itself, because
        // it's already in dpSmall.
        if (j == n / j)
            dpBig[i + 1][j] = 0;

        dpBig[i + 1][j] %= MOD;
    }
    for (int j = 1; j < sqrtN; j++) {
        // If we are adding a small number, we can add it to any sequence ending with a sufficiently small
        // big number, or any sequence at all ending with a small number.
        dpSmall[i + 1][j] = dpBig[i][j] + dpSmall[i][sqrtN - 1];
        dpSmall[i + 1][j] %= MOD;
    }
}

long answer = 0;
for (long x : dpSmall[k]) {
    answer += x;
    answer %= MOD;
}
for (long x : dpBig[k]) {
    answer += x;
    answer %= MOD;
}

out.println(answer);

Thanks for reading! Let me know if you have any questions or feedback for me.