Problem: Collect Coins in a Tree

There exists an undirected and unrooted tree with `n` nodes indexed from `0`
to `n - 1`. You are given an integer `n` and a 2D integer array edges of
length `n - 1`, where `edges[i] = [ai, bi]` indicates that there is an edge
between nodes `ai` and `bi` in the tree. You are also given an array `coins`
of size `n` where `coins[i]` can be either `0` or `1`, where `1` indicates the
presence of a coin in the vertex `i`.

Initially, you choose to start at any vertex in the tree. Then, you can
perform the following operations any number of times:

* Collect all the coins that are at a distance of at most `2` from the current vertex, or
* Move to any adjacent vertex in the tree.

Find the minimum number of edges you need to go through to collect all the
coins and go back to the initial vertex.

Note that if you pass an edge several times, you need to count it into the
answer several times.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class Solution:
  def collectTheCoins(self, coins: List[int], edges: List[List[int]]) -> int:
    n = len(coins)
    tree = [set() for _ in range(n)]
    leavesToBeRemoved = collections.deque()

    for u, v in edges:
      tree[u].add(v)
      tree[v].add(u)

    for u in range(n):
      while len(tree[u]) == 1 and coins[u] == 0:
        v = tree[u].pop()
        tree[v].remove(u)
        u = v
      if len(tree[u]) == 1:
        leavesToBeRemoved.append(u)

    for _ in range(2):
      for _ in range(len(leavesToBeRemoved)):
        u = leavesToBeRemoved.popleft()
        if tree[u]:
          v = tree[u].pop()
          tree[v].remove(u)
          if len(tree[v]) == 1:
            leavesToBeRemoved.append(v)

    return sum(len(children) for children in tree)


Java solution:
class Solution {
  public int collectTheCoins(int[] coins, int[][] edges) {
    final int n = coins.length;
    Set<Integer>[] tree = new Set[n];
    Deque<Integer> leavesToBeRemoved = new ArrayDeque<>();

    for (int i = 0; i < n; ++i)
      tree[i] = new HashSet<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      tree[u].add(v);
      tree[v].add(u);
    }

    for (int i = 0; i < n; ++i) {
      int u = i;
      // Remove the leaves that don't have coins.
      while (tree[u].size() == 1 && coins[u] == 0) {
        final int v = tree[u].iterator().next();
        tree[u].clear();
        tree[v].remove(u);
        u = v; // Walk up to its parent.
      }
      // After trimming leaves without coins, leaves with coins may satisfy
      // `leavesToBeRemoved`.
      if (tree[u].size() == 1)
        leavesToBeRemoved.offer(u);
    }

    // Remove each remaining leaf node and its parent. The remaining nodes are
    // the ones that must be visited.
    for (int i = 0; i < 2; ++i)
      for (int sz = leavesToBeRemoved.size(); sz > 0; --sz) {
        final int u = leavesToBeRemoved.poll();
        if (!tree[u].isEmpty()) {
          final int v = tree[u].iterator().next();
          tree[u].clear();
          tree[v].remove(u);
          if (tree[v].size() == 1)
            leavesToBeRemoved.offer(v);
        }
      }

    return Arrays.stream(tree).mapToInt(children -> children.size()).sum();
  }
}


C++ solution:
class Solution {
 public:
  int collectTheCoins(vector<int>& coins, vector<vector<int>>& edges) {
    const int n = coins.size();
    vector<unordered_set<int>> tree(n);
    queue<int> leavesToBeRemoved;

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      tree[u].insert(v);
      tree[v].insert(u);
    }

    for (int i = 0; i < n; ++i) {
      int u = i;
      // Remove the leaves that don't have coins.
      while (tree[u].size() == 1 && coins[u] == 0) {
        const int v = *tree[u].begin();
        tree[u].clear();
        tree[v].erase(u);
        u = v;  // Walk up to its parent.
      }
      // After trimming leaves without coins, leaves with coins may satisfy
      // `leavesToBeRemoved`.
      if (tree[u].size() == 1)
        leavesToBeRemoved.push(u);
    }

    // Remove each remaining leaf node and its parent. The remaining nodes are
    // the ones that must be visited.
    for (int i = 0; i < 2; ++i)
      for (int sz = leavesToBeRemoved.size(); sz > 0; --sz) {
        const int u = leavesToBeRemoved.front();
        leavesToBeRemoved.pop();
        if (!tree[u].empty()) {
          const int v = *tree[u].begin();
          tree[u].clear();
          tree[v].erase(u);
          if (tree[v].size() == 1)
            leavesToBeRemoved.push(v);
        }
      }

    return accumulate(tree.begin(), tree.end(), 0,
                      [](int acc, const unordered_set<int>& children) {
      return acc + children.size();
    });
  }
};