RAINDROPS - Editorial


Div-3 Contest
Div-2 Contest
Div-1 Contest

Author: Vishesh Saraswat
Tester: Istvan Nagy
Editorialist: Vishesh Saraswat




Graph, DFS, Prefix Sums


There is a tree with N nodes. Every second new raindrops fall on its leaves and old raindrops move 1 node closer to the root. The tree is shook M times, each shake resulting in all raindrops except the ones at the root falling off. Find the number of raindrops at the root after the $M$th shake.


We can see that each shake essentially resets the tree to the initial version. So, we can try to solve this problem for M independent time intervals and then add up the answer.

What we need now is an efficient method to calculate the answer for each interval. We know that a raindrop from a leaf which is at a distance of d nodes from the root, will take d+1 seconds to reach the root. After a raindrop reaches the root for the first time, more raindrops will follow from that leaf and we will get a raindrop from that leaf indefinitely. Building upon this, let’s now take D to be the distance of the leaf which is farthest from the root. Therefore after D seconds, we can be sure that the root is receiving X (the number of leaves) raindrops every second.

All we need to do now is to calculate how many raindrops the root receives for \leq D seconds after the tree is reset. To do so let’s define an array R such that R_i is the number of raindrops the root has after i seconds from a reset. We’ll also define an array L such that L_i is the number of leaves at a distance of i from the root. R_1 will be 0 because no drops can reach the root in the first second. Now, for all 1 \lt i \leq D we can calculate R_i as R_{i-1} + \sum_{j=1}^{i-1}L_j, the sum of raindrops that have already reached by i-1 seconds and the number of new raindrops that will reach this second.

So, we will calculate our final answer by going over all M intervals. Let’s say Y is the duration of current interval, we will add R_Y to the answer if Y \leq D, otherwise we will add R_D + (X \cdot (Y-D)) to the answer. X here is the number of total leaves in the tree.


Setter's / Editorialist's Solution
#include "bits/stdc++.h"
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
using ordered_set = tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>;

#define all(x) begin(x), end(x)
#define rall(x) rbegin(x), rend(x)
#define sz(x) (int)(x).size()

using ll = long long;
const int mod = 1e9+7;

vector<vector<int>> adj;
vector<int> L;
vector<int> dist;
ll X;

void dfs(int u, int p = -1) {
    for (int v : adj[u]) {
        if (v != p) {
            dist[v] = dist[u] + 1;
            dfs(v, u);
    if (sz(adj[u]) == 1 and u != 0) {

void solve(int tc) {
    int n, m;
    cin >> n >> m;
    X = 0;
    adj.clear(); L.clear(); dist.clear();
    adj.resize(n); L.resize(n); dist.resize(n);
    for (int i = 0; i < n-1; ++i) {
        int u, v;
        cin >> u >> v;
        --u, --v;
    vector<int> a(m+1);
    for (int i = 1; i <= m; ++i)
        cin >> a[i];
    ll D = n-1;
    while (L[D] == 0)
    vector<ll> R(D+1);
    ll cursum = 0;
    for (int i = 1; i <= D; ++i) {
        R[i] = R[i-1] + (cursum + L[i-1]);
        cursum += L[i-1];
    ll ans = 0;
    for (int i = 1; i <= m; ++i) {
        ll Y = a[i] - a[i-1];
        if (Y <= D)
            ans += R[Y];
            ans += R[D] + (Y - D) * X;
    cout << ans << '\n';

signed main() {
    int tc = 1;
    cin >> tc;
    for (int i = 1; i <= tc; ++i) solve(i);
    return 0;
Tester's Solution
#include <iostream>
#include <algorithm>
#include <string>
#include <cassert>
#include <vector>
#include <numeric>
using namespace std;

#ifdef HOME
#define NOMINMAX
#include <windows.h>

long long readInt(long long l, long long r, char endd) {
	long long x = 0;
	int cnt = 0;
	int fi = -1;
	bool is_neg = false;
	while (true) {
		char g = getchar();
		if (g == '-') {
			assert(fi == -1);
			is_neg = true;
		if ('0' <= g && g <= '9') {
			x *= 10;
			x += g - '0';
			if (cnt == 0) {
				fi = g - '0';
			assert(fi != 0 || cnt == 1);
			assert(fi != 0 || is_neg == false);

			assert(!(cnt > 19 || (cnt == 19 && fi > 1)));
		else if (g == endd) {
			assert(cnt > 0);
			if (is_neg) {
				x = -x;
			assert(l <= x && x <= r);
			return x;
		else {

string readString(int l, int r, char endd) {
	string ret = "";
	int cnt = 0;
	while (true) {
		char g = getchar();
		assert(g != -1);
		if (g == endd) {
		ret += g;
	assert(l <= cnt && cnt <= r);
	return ret;
long long readIntSp(long long l, long long r) {
	return readInt(l, r, ' ');
long long readIntLn(long long l, long long r) {
	return readInt(l, r, '\n');
string readStringLn(int l, int r) {
	return readString(l, r, '\n');
string readStringSp(int l, int r) {
	return readString(l, r, ' ');

int main() {
#ifdef HOME
	if (IsDebuggerPresent())
		freopen("../in.txt", "rb", stdin);
		freopen("../out.txt", "wb", stdout);
	int T = readIntLn(1, 1000);
	uint32_t sumN = 0;
	for (int tc = 0; tc < T; ++tc)
		uint32_t N = readIntSp(2, 100'000);
		sumN += N;
		uint32_t M = readIntLn(1, 100'000);
		vector<vector<uint32_t>> neighb(N);
		for (uint32_t i = 0; i < N - 1; ++i)
			uint32_t u = readIntSp(1, N);
			uint32_t v = readIntLn(1, N);
			assert(u != v);
			--u; --v;
		vector<uint32_t> A(M);
		vector<uint32_t> B;
		uint32_t actr = 0;
		uint32_t prevA = 0;
		for (auto&& ai : A)
			if (actr == M)
				ai = readIntLn(1, 1'000'000'000);
				ai = readIntSp(1, 1'000'000'000);
			assert(ai > prevA);
			prevA = ai;
		vector<int> d;
		vector<int> nodes({0});
		vector<bool> used(N);
		used[0] = true;
		for (uint32_t i = 0; i < nodes.size(); ++i)
			auto actNode = nodes[i];
			for (auto cand : neighb[actNode])
				used[cand] = true;
				d[cand] = d[actNode] + 1;
		assert(nodes.size() == N);
		uint32_t maxDepth = 0;
		uint64_t leafCount = 0;
		for (uint32_t i = 1; i < N; ++i)
			if (neighb[i].size() == 1)
				maxDepth = std::max<uint32_t>(maxDepth, d[i]);
		vector<uint32_t> vRainDrops(maxDepth + 1);
		for (uint32_t i = 1; i < N; ++i)
			if (neighb[i].size() == 1)
		uint32_t dropCtr = 0;
		for (uint32_t i = 1; i < vRainDrops.size(); ++i)
			dropCtr += vRainDrops[i];
			vRainDrops[i] = vRainDrops[i-1] + dropCtr;
		uint64_t res = 0;
		uint32_t prev = 0;
		for (auto ai : A)
			uint64_t dist = ai - prev;
			if (dist < vRainDrops.size())
				res += vRainDrops[dist];
				res += vRainDrops[maxDepth] + (dist - maxDepth) * leafCount;
			prev = ai;
		printf("%llu\n", res);
	assert(sumN <= 500'000);
	assert(getchar() == -1);

The event is a part of our annual tech symposium Exun 2021-22, powered by Athena Education.


Didnt notice Ai is upto 1e9, thought it should be <=n . :frowning_face: :confounded:

Can’t we call a root node as a leaf node, isn’t it possible ?

Sample Input:

4 1
1 2
2 3
2 4


      /   \
     3     4  

according to your statement why cannot at time t=0 a water drop fall at node 1 (root node) anyway it is a leaf node.

If we can call the root node as leaf node, then the correct answer for the given test case would be 146.

But the Accepted Solutions and the Editorialist’s solution gives, 96. considering 1 as non leaf node.


No, root can be considered as a leaf only if number of nodes is 1.
Here in the constraints its mentioned n>=2 , so rain drop can’t fall on the root as it cannot be considered as a leaf.


In the example you gave, only if it is a regular tree the node 1 can be considered a leaf node. But we’re talking about a rooted tree here. The definition of a leaf node changes in a rooted tree to a non-root node with 1 neighbour. So the root node cannot be a leaf node no matter what.
Hope this isn’t too confusing :sweat_smile:

1 Like

i’m getting run time error in one subtask and all other subtask are AC.
I’m not able to find out the wrong can some one help.

from collections import defaultdict

def solve(d, i, p, depth, aa):
    for x in d[i]:
        if x != p:
            solve(d, x, i, depth + 1, aa)
    if len(d[i]) == 1 and i != 0 and depth<len(aa):
        aa[depth] += 1

for _ in range(int(input())):
    n, m = map(int, input().split())
    d = defaultdict(list)
    for i in range(n - 1):
        u, v = map(int, input().split())
        d[u - 1].append(v - 1)
        d[v - 1].append(u - 1)
    am = list(map(int, input().split()))
    aa = [0 for i in range(n + 1)]
    solve(d, 0, -1, 0, aa)
    ans = 0
    if len(am)>=1:
        q = [am[0]]
    for i in range(1, len(am)):
        q.append(am[i] - am[i - 1])
    aaa = []
    for i in range(0, n + 1):
        aaa.append(aa[i] * i)

    for i in range(1, n + 1):
        aa[i] += aa[i - 1]
        aaa[i] += aaa[i - 1]

    for i in range(0, m):
        ans += q[i] * aa[min(n, q[i])] - aaa[min(n, q[i])]

i came up with exact same logic but messed up in implementation :no_mouth:

1 Like

BFS + prefix sum solution

1 Like

This is a type of problem where implementing it is bit harder then framing the logic.

Here is a thought process :
First of all I stored all the leaf nodes in a set. 1 cannot be a leaf node so exclude that.
Here leaf nodes are the ones which has only one edge attached to it.

So what I did here is that I traversed the entire tree in a dfs manner calculating the height of each node.
Whenever I came across a leaf node I had its distance from the root node. Let h be that distance.

Now after every h+1 second the root node will receive a droplet from this leaf node so I incremented Time[h+1].
Initially , Time[i] indicates the number of droplets that I will receive exactly at ith second.

Now I perform the prefix operation on the Time array i.e. Time[i] += Time[i-1].
After this Time[i] indicates the total number of droplets collected at ith second.
// This is not true.

But here is the trick once a single droplet has completed the path to root then after each second the root will receive a droplet from that leaf.
So we perform another prefix operation and now Time[i] indicates the total number of droplets collected after at ith second.

Here is the Submission Link

Can you please tell me how to remove the tle!

You are storing the distances in a set and then fetching it one by one for every A[i] , hence its linear time.
Precompute the values and store it in prefix sums, by which you can do the same thing in O(1).

My Submission

Also I think using BFS instead of DFS is much better here because of multiple testcases and hence have to clear the entire the adjacency list and all the global vectors/variables in the beginning of every testcase which is very intimidating.

I wasted about 20min in the contest because this was not working.

for(auto v:adj)

And had to go with this:-

for (int i = 0; i < MAXN; ++i)