公式

G - Mode in the Subtree 解説 by en_translator


At a glance, the problem seems solvable with the merging trick (“merging into the larger cluster” strategy) with an unordered_map (hash map) in \(\mathrm{O}(N \log N)\) time. However, the problem has a very large constraints, \(N \leq 2.5 \times 10^6\), while unordered_map has a very heavy constant factor for both time and spacial complexity, so the problem cannot be probably solved by naively adopting the trick.

The writer’s solution utilizes an algorithm called DSU (Disjoint Set Union) on Tree to eliminate the disadvantage of the trick, that is, a bad constant factor due to an overly functional data structure.

Let us introduce the outline of DSU on Tree. A DSU on Tree is an algorithm that computes information regarding a subtree rooted at some vertex by cleverly incorporating the results for the children.
For each vertex \(c\), we want to obtain a table that stores information about vertices contained in its subtree. Using the merging trick is a possible naive approach, but direct implementation, especially if it uses a heavy data structure like unordered_map, results in a quite huge constant factor.
Instead, for each vertex, define the child whose subtree has the most vertices as the “heavy child,” and the others as “light children.” We reuse the table for a heavy child without discarding it, and incorporate information for light children only when needed. This way, the time complexity can be reduced, allowing us to solve the problem without using a data structure with a heavy constant factor.

We will now describe the detailed algorithm of DSU on Tree. First, obtain the Heavy-Light Decomposition (HLD) of the tree. (Regarding the heavy-light decomposition algorithm, please refer to ABC269-Ex editorial.) Classify each edge into a heavy edge or a light edge. Then, perform a DFS represented by the following pseudocode:

# We aimt at managing information associated with each vertex in a table
# add(c): Insert information of vertex c into the table
# query(c): Respond to querie against a table containing infromation about the entire subtree rooted at vertex c
# reset(): Reset the table

# c is vertex number
def dfs(c):
  # First, process light children
  for d in (light children of c):
    dfs(d)
    # Reset the table for a light child
    reset()

  # Reuse the table for the subtree rooted at the heavy child
  if c is not leaf:
    dfs(heavy child of c)
  
  # Incorporate the information of the descendants of the light children into the table
  for d in (light children of c):
    for n in (descendants of c):
      add(n)

  # Incorporte the information of c itself into the table
  add(c)

  # Answer the query
  query(c)
  
  return

Let us analyze the complexity of this algorithm. For a vertex n, the number of times add function is called is (the number of light edges contained in the path from the root to vertex n) + 1. This is because the tables are inherited from heavy children during the DFS. Therefore, throughout the algorithm:

  • add function is called \(\mathrm{O}(N \log N)\) times;
  • query function is called \(N\) times; and
  • reset function is called \(\mathrm{O}(N)\).

Hence, as long as add, query, and reset functions operate fast enough, the time complexity of the entire algorithm is \(\mathrm{O}(N \log N)\).

In our problem, it is sufficient to store the following information:

  • cnt[x]: the number of vertices with color x among the vertices currently stored in the table
  • num[t]: the number of colors that occur exactly t times among the colors currently stored in the table
  • mx: the maximum frequency

Since color numbers do not exceed \(N\), all of them can be managed in an array. To add vertex v in add(v), let x = c_v, increment cnt[x] by one, and update num and mx accordingly. The answer to query(v) is simply (mx, num[mx]). For reset(), we do not need to clear the entire array; instead, we may record the updated colors and frequencies, and set only them back to \(0\).

The solution using the merge trick was inefficient in terms of both time and spacial complexity, due to the bad constant factor of unordered_map. On the other hand, DSU on Tree realizes the idea of “merging the smaller into the larger” using arrays. As in our case, if all the information can be processed on arrays, we can say that DSU on Tree is more practical.

The problem can be solved by appropriately implementing the solution above. The time complexity is \(\mathrm{O}(N \log N)\), and the spacial complexity is \(\mathrm{O}(N)\), which are both small enough.

  • Sample code (Python, codon): if a somewhat small constant factor is required, using Euler Tree together with HLD may help improving the factor.
N, state, M, F = [int(x) for x in input().split()]
Q = [int(x) for x in input().split()]
D = [int(x) for x in input().split()]

p = [-1] * N
g = [list() for _ in range(N)]
for i in range(2, N + 1):
    if i <= M:
        p[i - 1] = Q[i - 2] - 1
    else:
        p[i - 1] = state % (i - 1)
        state = (state * 1103515245 + 12345) & ((1 << 31) - 1)
    g[p[i - 1]].append(i - 1)

C = [-1] * N
for i in range(1, N + 1):
    if i <= M:
        C[i - 1] = D[i - 1] - 1
    else:
        C[i - 1] = state % F
        state = (state * 1103515245 + 12345) & ((1 << 31) - 1)

sub = [0] * N
for c in reversed(range(N)):
    sub[c] = 1
    for j in range(len(g[c])):
        d = g[c][j]
        sub[c] += sub[d]
        if sub[d] > sub[g[c][0]]:
            g[c][j], g[c][0] = g[c][0], g[c][j]

idx, euler, down, up = 0, [0] * N, [0] * N, [0] * N
st = [~0, 0]
while len(st) > 0:
    c = st.pop()
    if c >= 0:
        euler[idx] = c
        down[c] = idx
        idx += 1
        for d in reversed(g[c]):
            st.append(~d)
            st.append(d)
    else:
        up[~c] = idx

mx, ans, cnt, freq = 0, 0, [0] * N, [0] * (N + 1)


def add(c):
    global mx
    col = C[c]
    freq[cnt[col]] -= 1
    cnt[col] += 1
    freq[cnt[col]] += 1
    mx = max(mx, cnt[col])


def query(c):
    global ans, mx
    ans += (mx ^ (c + 1)) * (freq[mx] ^ (c + 1))
    ans %= 998244353


def reset(c):
    global mx
    for i in range(down[c], up[c]):
        col = C[euler[i]]
        cnt[col] = 0
    for i in range(mx + 1):
        freq[i] = 0
    mx = 0


def dsu(c):
    global mx, ans
    for i in range(1, len(g[c])):
        dsu(g[c][i])
        reset(g[c][i])
    if len(g[c]):
        dsu(g[c][0])
        for i in range(up[g[c][0]], up[c]):
            add(euler[i])
    add(c)
    query(c)


dsu(0)
print(ans)

投稿日時:
最終更新: