This problem features matroid intersection.
While solving matroid intersection problem in general requires a complex algorithm, for special cases called linear matroid intersection there is a simple randomized algorithm that does not require complex implementation.

Definition of matroid

A finite set \(E\) and a family of its subsets \(I\) are called a matroid \((E, I)\) if and only if they satisfy all of the following conditions:

• (1): \(I\) contains the empty set.
• (2): If \(I\) contains \(X\), then it also contains every subset \(Z\) of \(X\).
• (3): If \(X, Y \in I\) and \(|X| < |Y|\), then there exists some \(e \in Y \setminus X\) such that \(X \cup \lbrace e \rbrace \in I\).

Additionally, we define the following terms:

• Independent set: A set contained in \(I\).
• Dependent set: A subset of \(E\) that is not in \(I\).
• Basis: A maximal independent set.
• For \(X \subseteq E\), a basis of \(X\): A maximal independent set contained in \(X\).

The third condition in the definition of a matroid is equivalent to the following condition (proof omitted):

• (3’): For all \(X \subseteq E\), all bases of \(X\) have the same size.

In particular, considering the case \(X = E\), we see that all bases of a matroid have the same size.

Examples of Matroids

Let’s look at some concrete examples of matroids.

Uniform Matroid and Partition Matroid

Given a finite set \(E\) and a non-negative integer \(r\), define a family of subsets \(I\) as follows:

\[I = \lbrace X \subseteq E \mid \vert X \vert \leq r \rbrace\]

Then, \((E, I)\) forms a matroid, which is called a uniform matroid.
Moreover, the direct sum of uniform matroids also forms a matroid, known as a partition matroid. Specifically, given a partition of \(E\), denoted by \(\lbrace E_1, E_2, \dots, E_k \rbrace\), and non-negative integers \(r_1, r_2, \dots, r_k\), define \(I\) as follows:

\[I = \lbrace X \subseteq E \mid \forall i \in \lbrace 1,2,\dots, k \rbrace, \vert X \cap E_i \vert \leq r_i \rbrace\]

Then, \((E, I)\) is a matroid. (The proof is straightforward for all conditions (1), (2), and (3).)

Linear Matroid

Consider a matrix \(A\), and let \(E\) be the set of columns of \(A\). Define \(I\) as follows:

\[I = \lbrace X \subseteq E \mid \text{the columns in } X \text{ are linearly independent} \rbrace\]

Then, \((E, I)\) forms a matroid, which is called a linear matroid. (It can be proven since (1), (2), and (3’) all follow directly from linear algebra.)

Cycle Matroid (Graphical Matroid)

Let \(E\) be the set of edges of a graph \(G\), and let \(V\) be the set of vertices of \(G\). Define \(I\) as follows:

\[I = \lbrace X \subseteq E \mid \text{the graph } (V, X) \text{ contains no cycles} \rbrace\]

Then, \((E, I)\) forms a matroid, which is called a cycle matroid (also known as a graphical matroid). (The proof follows easily from competitive programming knowledge, as conditions (1), (2), and (3’) hold naturally.)

Matroid Intersection

A frequently encountered topic in problems related to matroids is matroid intersection.
Given two matroids \((E, I_1)\) and \((E, I_2)\), their intersection \(I_1 \cap I_2\) is defined as follows:

\[I_1 \cap I_2 = \lbrace X \mid X \in I_1, X \in I_2 \rbrace\]

The problem of finding \(X \in I_1 \cap I_2\) with the maximum \(\vert X \vert\) is called the matroid intersection problem.

A well-known example of a problem that can be described using matroid intersection is the maximum matching in a bipartite graph. Consider a bipartite graph with partite sets \(U\) and \(V\) and an edge set \(E\). The matching condition on \(U\) requires that each \(u \in U\) has at most one adjacent edge selected, which corresponds to a partition matroid. Similarly, the matching condition on \(V\) can also be expressed as a partition matroid. Therefore, bipartite graph matching can be seen as the intersection of two partition matroids. Many other problems can also be formulated as matroid intersection.

It is known that the matroid intersection problem can be solved in polynomial time (approximately \(\mathrm{O}(|E|^2 r)\), where \(r\) is the size of the solution set), but the algorithm is complex and is omitted here. For more details, please refer to the following resources (in Japanese):

• Lecture notes by Yoshio Okamoto, Implementation example by hitonanode

Linear Matroid Intersection

A special case of matroid intersection is the linear matroid intersection problem.
It is known that if we only seek the size of the solution, the linear matroid intersection problem can be solved more concisely than the general matroid intersection problem.

We first formulate the problem. Let \(A_1\) and \(A_2\) be matrices with the same number of columns, and let \(E\) be the set of columns. Then, we can define linear matroids \((E, I_1)\) and \((E, I_2)\) corresponding to \(A_1\) and \(A_2\), respectively. The intersection \(I_1 \cap I_2\) is called the linear matroid intersection.

The following proposition holds for the linear matroid intersection problem:

Define polynomial matrices \(D\) and \(M\) using variables \(x_1, x_2, \dots, x_{\vert E \vert}\) as follows:

\(D = \mathrm{diag}(x_e)\) (an \(E \times E\) diagonal matrix with \(x_e\) on the diagonal),

\[M = A_1 D A_2^{T}.\]

Then, the size of the solution to the linear matroid intersection problem is equal to \(\mathrm{rank}(M)\).

For a proof, refer to Taihei Oshiro’s materials.

By the proposition above, solving the linear matroid intersection problem reduces to computing the rank of a polynomial matrix. Directly computing the rank of a polynomial matrix is difficult, but it can be found with high probability by substituting random values for the variables and computing the rank of the resulting matrix modulo \(p\). (The validity of this approach is guaranteed by the Schwarz-Zippel lemma.) Thus, letting \(n\) be the larger row count of \(A_1\) and \(A_2\), and \(m\) be the smaller row count, we can compute the solution size of the linear matroid intersection problem in approximately \(\mathrm{O}(nm^2)\) time.

Concrete Example of Linear Matroid Intersection (Edmonds Matrix)

As a concrete example of the linear matroid intersection, let’s consider the maximum matching problem in a bipartite graph and the Edmonds matrix.
Let the sets of partitions be \(U\) and \(V\), and the set of edges be \(E\). The Edmonds matrix \(N\) for the bipartite graph is defined as follows:

\[N_{u,v} = \begin{cases} x_{u,v}& &\text{if } (u,v) \in E \\ 0 & &\text{otherwise.} \end{cases} \]

In this case, the size of the maximum matching in the bipartite graph is equal to \(\mathrm{rank}(N)\). This fact can be proven without using linear matroid intersection, but here we will prove it using linear matroid intersection.

As mentioned earlier, the maximum matching in a bipartite graph can be expressed as the intersection of two partition matroids. By further transforming it as follows, we can reduce it to a linear matroid intersection. Let \(\mathbf{e}_i\) represent a column vector where the \(i\)-th component is 1, and the other components are 0. We define the matrices \(A_1\) and \(A_2\) such that the column vector in \(A_1\) corresponding to an edge \(e=(u,v)\) is \(\mathbf{e}_u\) and that in \(A_2\) is \(\mathbf{e}_v\). By choosing \(A_1\) and \(A_2\) in this way, we can express the matching in the bipartite graph as the intersection of the linear matroids \(I_1\) and \(I_2\) corresponding to \(A_1\) and \(A_2\), respectively.

Now, calculating the matrix \(M\), we get:

\[ \begin{aligned} M &= A_1 D A_2^{T} \\ &= \sum_{e=(u,v) \in E} \mathbf{e}_u \mathbf{e}_v^T x_{u,v} \end{aligned} \]

Here, \(\mathbf{e}_u \mathbf{e}_v^T\) is a matrix where only the \(uv\)-th component is 1, and all other components are 0. Therefore, we can confirm that \(M\) matches the Edmonds matrix \(N\), and from this, we can conclude that the size of the maximum matching in the bipartite graph is equal to \(\mathrm{rank}(N)\).

Now that we’ve covered the content related to linear matroid intersection, let’s think about how to solve the problem at hand using linear matroid intersection.

Solution to the Problem

As the title suggest, the problem features colorful spanning tree. A colorful spanning tree refers to a spanning tree defined as in the problem statement (although academically, it seems to refer to the case where \(A_i = 1\)). The problem of finding a colorful spanning tree is a well-known example of matroid intersection.

Let’s first consider the simpler case of determining the existence of a colorful spanning tree without the constraints on \((L, R)\). Let the graph be \((V, E)\), and let \(E_c\) be the set of edges of color \(c\). If we define

\[I_1 = \lbrace X \subseteq E \mid \forall c \in \lbrace 1,2,\dots, C \rbrace, \vert X \cap E_c \vert \leq A_c \rbrace,\]

\[I_2 = \lbrace X \subseteq E \mid \text{Graph }(V, X) \text{ does not contain a cycle} \rbrace,\]

then \((E, I_1)\) and \((E, I_2)\) are a partition matroid and a cycle matroid, respectively. The condition for an edge set \(X\) to form a colorful spanning tree is that \(X \in I_1 \cap I_2\) and \(\vert X \vert = N-1\). Thus, combining the fact that the size of a basis of \(I_2\) is \(N-1\), we can conclude that the problem of determining the existence of a colorful spanning tree is equivalent to determining whether the largest set in the intersection \(I_1 \cap I_2\) has size \(N-1\).

Therefore, we have found that this problem can be solved in polynomial time using the matroid intersection algorithm. Let’s further transform the problem to reduce it to a linear matroid intersection. In other words, we will transform both the partition matroid and the cycle matroid into linear matroids corresponding to a certain matrix.
Let \(A_1\) and \(A_2\) be matrices with modulo-\(p\) integers corresponding to \(I_1\) and \(I_2\), respectively (where \(p\) is a sufficiently large prime number). We need to carefully design \(A_1\) and \(A_2\) to establish an appropriate correspondence.
Designing \(A_2\) is relatively simple: we can set the number of rows to \(N\), and the column vector corresponding to each edge \((u, v)\) to the one whose \(u\)-th component is \(1\) and the \(v\)-th is \(-1\). The validity can be proved from the fact that the minimal dependent set of column vectors forms a cycle.
Designing \(A_1\) is less obvious; we want a structure where at most \(A_c\) column vectors of color \(c\) can be chosen. One possible design is as follows:

• Let \(\sum_{c=1}^C A_c = K\). \(A_1\) will be a matrix with \(K\) rows and \(\vert E \vert\) columns.
• Assign the first \(A_1\) rows to correspond to color \(1\), the next \(A_2\) rows to color \(2\), and so on.
• The column vector corresponding to an edge of color \(c\) is defined as follows:
  • Suppose that the current vector is \(n\)-th such vector. Then, for the rows corresponding to color \(c\), the first row will contain \(n^0\), the second row will contain \(n^1\), and so on up to the \(A_c\)-th row, which will contain \(n^{A_c-1}\). All other rows will be filled with 0.

The validity of \(A_1\) follows from the Vandermonde determinant being non-zero.

Now, with the method described above, we have reduced the matroid intersection to a linear matroid intersection. Therefore, for the case where there are no constraints on \(L\) and \(R\), the problem can be solved by applying the randomized algorithm for linear matroid intersection.

Let us now consider the original problem, i.e., the problem of counting the number of valid pairs \((L, R)\) that satisfy the condition in the problem statement. Focusing on the matrix \(M = A_1 D A_2^T\) that appears in the randomized algorithm, we observe that the edges of color \(c\) only affect the rows corresponding to color \(c\). Let \(S_c = \sum_{i=1}^{c-1} A_i\), and we find that for \((L, R)\) to satisfy the condition in the problem statement, it is equivalent to the rank of the submatrix formed by rows \(S_L+1\) through \(S_R\) of \(M\) being \(N-1\). Therefore, if we can solve the following problem, we can determine whether the condition are satisfied for each \((L, R)\).

Solve the following problem for \(c=1,2,\dots,C\):

• Start from row \(S_c+1\) and proceed downward. What is the first row where the rank of the submatrix formed by the rows seen so far becomes \(N-1\)?

This problem can be solved for each \(c=1,2,\dots,C\) using an algorithm that adds row vectors while maintaining the basis, with a time complexity of \(\mathrm{O}(N^2 C K)\). Furthermore, using the technique introduced in ABC223H, the time complexity can be reduced to \(\mathrm{O}(N^2 K)\). (A detailed explanation is omitted.)

By appropriately implementing these observations, we can solve the problem in \(\mathrm{O}(N^2 K)\) time with \(\sum_{c=1}^C A_c = K\), which is sufficiently fast. Since the constraints are relatively relaxed, an \(\mathrm{O}(N^2 C K)\) solution will also pass as long as the constant factor is fine.

• As an alternative solution, let’s consider solving this problem using a general matroid intersection algorithm. In fact, the following fact holds:
  • For each color \(c\), if we take a maximal sub-forest and remove the edges not belonging to it, the answer to the problem remains unchanged.
• Using this fact, we can reduce the size of the edge set \(E\) at an order level, and by using a general matroid intersection algorithm with a sliding window technique, the algorithm will run in \(\mathrm{O}(N^3 C^2)\) time. This will likely result in a TLE (Time Limit Exceeded) as it is, but with constant factor speedups and some effective pruning, we presume that the solution may pass.

Bonus

1. By applying the current algorithm, we can solve the weighted linear matroid intersection problem, that is, the problem of finding the maximum value of \(\sum_{e \in X} w_e\) for \(X \in I_1 \cap I_2\), in \(\mathrm{O}(W n^4)\) time, where \(n\) is the number of rows in the matrix and \(W\) is the maximum weight. Let’s think about this.

2. This problem (spoiler alert) is actually inspired by the matroid intersection problem, and the answer can be expressed as the determinant of a certain matrix. Let’s try interpreting this problem using matroid intersection.