[Idea] Using HLD to reduce memory

Правка en10, от maomao90, 2024-04-07 12:31:56

Recently when I was doing Universal Cup Round 5, I got stuck on a tree problem A as I realized that my solution required way too much memory. However, after the contest, I realized that there was a way that I could reduce a lot of memory using HLD. So here I am with my idea...

Structure of Tree DP

Most tree DP problems follow the following structure.

struct S {
    // return value of DP
};
S init(int u) {
    // initialise the base state of dp[u]
}
S merge(S left, S right) {
    // returns the new dp state where old state is left and transition using right
}
S dp(int u, int p) {
    S res = init(u);
    for (int v : adj[u]) {
        if (v == p) continue;
        res = merge(res, dp(v, u));
    }
    return res;
}
int main() {
    dp(1, -1);
}

An example of a tree DP using this structure is maximum independent set (MIS).

Code

Suppose struct $$$S$$$ requires $$$|S|$$$ bytes and our tree has N vertices. Then this naive implementation of tree DP requires $$$O(N\cdot |S|)$$$ memory as res of the parent is stored in the recursion stack as we recurse down to the leaves. This is fine for many problems as most of the time, $$$|S| = O(1)$$$, however in the case of the above question, $$$|S| = 25^2\cdot 24$$$ bytes and $$$N = 10^5$$$, which will require around $$$1.5$$$ gigabytes of memory, which is too much to pass the memory limit of $$$256$$$ megabytes. Below, I will show a way to use only $$$O(N + |S|\log N)$$$ memory.

Optimization

We try to make use of the idea of HLD and visit the vertex with the largest subtree size first.

struct S {
    // return value of DP
};
S init(int u) {
    // initialise the base state of dp[u]
}
S merge(S left, S right) {
    // returns the new dp state where old state is left and transition using right
}
int sub[MAXN];
void getSub(int u, int p) {
    sub[u] = 1;
    pair<int, int> heavy = {-1, -1};
    for (int i = 0; i < adj[u].size(); i++) {
        int v = adj[u][i];
        if (v == p) continue;
        getSub(v, u);
        sub[u] += sub[v];
        heavy = max(heavy, {sub[v], i});
    }
    // make the vertex with the largest subtree size the first
    if (heavy.first != -1) {
        swap(adj[u][0], adj[u][heavy.second]);
    }
}
S dp(int u, int p) {
    // do not initialize yet
    S res;
    bool hasInit = false;
    for (int v : adj[u]) {
        if (v == p) continue;
        S tmp = dp(v, u);
        if (!hasInit) {
            res = init();
            hasInit = true;
        }
        res = merge(res, tmp);
    }
    if (!hasInit) {
        res = init();
        hasInit = true;
    }
    return res;
}
int main() {
    getSub(1, -1);
    dp(1, -1);
}

If we analyze the memory complexity properly, we will realize that it becomes $$$O(N + |S|\log N)$$$. The $$$O(N)$$$ comes from storing the subtree sizes, and the $$$O(|S|\log N)$$$ comes from the DP itself.

Proof

The two main changes from our naive DP are that we initialize our res only after we visit the first child, and we visit the child with the largest subtree size first. Recalling the definitions used in HLD, we define an edge $$$p_u \rightarrow u$$$ to be a heavy edge if $$$u$$$ is the child with the largest subtree size among all children of $$$p_u$$$. We will call all other non-heavy edges light edges. It is well known that the path from the root to any vertex of a tree will involve traversing many heavy edges but at most $$$O(\log N)$$$ light edges.

Making use of this idea, if we visit heavy edges first, res of the parent of heavy edges will not be stored in the recursion stack since we only initialize our res after visiting the first edge, hence only res of the parent of light edges will be stored in the recursion stack. Then since there is at most $$$O(\log N)$$$ light edges, the memory complexity is $$$O(|S|\log N)$$$.

My solution to the above problem

Conclusion

I have not seen this idea anywhere before and only came up with it recently during the Universal Cup. Please tell me if it is a well-known technique or if there are some flaws in my analysis. Hope that this will be helpful to everyone!

Теги hld, memory

История

 
 
 
 
Правки
 
 
  Rev. Язык Кто Когда Δ Комментарий
en10 Английский maomao90 2024-04-07 12:31:56 126 Tiny change: '.\n\n[cut]\n\n# Opti' -> '.\n\n[cut] \n\n# Opti'
en9 Английский maomao90 2023-03-04 03:27:10 27 Tiny change: '\n ' -> '\n sub[u] += sub[v];\n '
en8 Английский maomao90 2023-03-03 13:23:19 23 Tiny change: '\n ' -> '\n getSub(v, u);\n '
en7 Английский maomao90 2023-03-03 09:50:46 21
en6 Английский maomao90 2023-03-03 09:32:29 36
en5 Английский maomao90 2023-03-03 07:51:24 1506 (published)
en4 Английский maomao90 2023-03-03 07:12:35 132
en3 Английский maomao90 2023-03-03 07:11:07 3083 Tiny change: 'uires $O(N|S|)$ memo' -> 'uires $O(N\cdot |S|)$ memo'
en2 Английский maomao90 2023-03-02 19:01:59 153
en1 Английский maomao90 2023-03-02 19:00:27 222 Initial revision (saved to drafts)