Open In App

Count Number of Nodes With Value One in Undirected Tree

Last Updated : 31 May, 2024
Comments
Improve
Suggest changes
Like Article
Like
Report

Given an undirected connected tree with n nodes labeled from 1 to n, and an integer array of queries[]. Initially, all nodes have a value of 0. For each query in the array, you need to flip the values of all nodes in the subtree of the node with the corresponding label. The parent of a node with label v is the node with label floor(v/2), and the root has label 1.

Note: Return the total number of nodes with a value of 1 after processing all the queries.

Example:



Example
Example


Input: n = 5 , queries = [1,2,5]
Output: 3
Explanation: The diagram above shows the tree structure and its status after performing the queries. The green node represents the value 0, and the grey node represents the value 1. After processing the queries, there are three grey nodes (nodes with value 1): 1, 3, and 5.

Input: n = 3, queries = [2,3,3]
Output: 1

Approach:

The main idea is that flipping a node's value twice brings it back to its original state. Therefore, we can simply track the number of times each node is flipped by the queries. If a node is flipped an odd number of times, its final value will be 1; otherwise, it will be 0.

We can use a depth-first search (DFS) to traverse the tree and count the number of flips for each node. For each node, we check if it needs to be flipped based on the query information. Then, we recursively explore its left and right subtrees, keeping track of the cumulative flips. Finally, we return the sum of the node's value and the values of its subtrees, which represents the total number of nodes with value 1 in its subtree.

Steps-by-step approach:

  • Initialize:
    • Create an array f of size n + 1, where n is the number of nodes in the tree.
    • Initialize all elements in f to 0.
    • Create a vector queries to store the queries.
  • Process Queries:
    • For each query in queries:
      • Flip the value of f[query] (0 becomes 1, 1 becomes 0).
    • Perform Depth-First Search (DFS):
  • Define a recursive function dfs(u, v):
    • If u is greater than n, return 0.
    • Flip the value of v based on f[u].
    • Recursively call dfs(u * 2, v) to explore the left subtree.
    • Recursively call dfs(u * 2 + 1, v) to explore the right subtree.
    • Return the sum of v (the value of the current node) and the results of the left and right subtrees.
  • Start DFS:
    • Call dfs(1, 0) to start the DFS from the root node (node 1) with an initial value of 0.
  • Return Result:
    • The result of dfs(1, 0) is the total number of nodes with value 1 in the tree. Return this value.

Below is the implementation of the above approach:

C++
#include <bits/stdc++.h>
using namespace std;

int numberOfNodes(int n, vector<int>& queries)
{
    // f[i] = 1 means we need to flip all values in the
    // subtree of the node i
    vector<int> f(n + 1);
    // if we flip the node even times, the value would be
    // same as the original value e.g. 0 (original value) ->
    // 1 -> 0 if we flip the node odd times, the value would
    // be the the opposite of the original value e.g.  0
    // (original value) -> 1 -> 0 -> 1 therefore, we can
    // first process the queries to obtain the final flips
    for (auto q : queries)
        f[q] ^= 1;

    function<int(int, int)> dfs = [&](int u, int v) {
        // u is the current node label
        // if u is greater than n, then return 0
        if (u > n)
            return 0;
        // do we need to flip the node u?
        // we flip the value if f[u] is 1
        v ^= f[u];
        // the result would be the value of u, i.e. v
        // plus the result of the left subtree, i.e dfs(u *
        // 2, v) plus the result of the right subtree, i.e.
        // dfs(u * 2 + 1, v)
        return v + dfs(u * 2, v) + dfs(u * 2 + 1, v);
    };
    // we start from node 1 with inital value 0
    return dfs(1, 0);
}

// Driver code
int main()
{
    int n = 7;
    vector<int> queries = { 3, 1, 4, 1, 5, 9, 2, 6 };

    int result = numberOfNodes(n, queries);
    cout << result << endl;

    return 0;
}
Java
import java.util.List;
import java.util.function.BiFunction;

public class NumberOfNodes {
    public static int numberOfNodes(int n,
                                    List<Integer> queries)
    {
        // f[i] = 1 means we need to flip all values in the
        // subtree of the node i
        int[] f = new int[n + 1];

        // If we flip the node even times, the value would
        // be same as the original value If we flip the node
        // odd times, the value would be the opposite of the
        // original value Therefore, we can first process
        // the queries to obtain the final flips
        for (int q : queries) {
            if (q <= n) {
                f[q] ^= 1;
            }
        }

        // Define the dfs function using a lambda expression
        BiFunction<Integer, Integer, Integer> dfs
            = new BiFunction<>() {
                  @Override
                  public Integer apply(Integer u, Integer v)
                  {
                      // u is the current node label
                      // if u is greater than n, then return
                      // 0
                      if (u > n)
                          return 0;
                      // do we need to flip the node u?
                      // we flip the value if f[u] is 1
                      v ^= f[u];
                      // the result would be the value of u,
                      // i.e. v plus the result of the left
                      // subtree, i.e dfs(u * 2, v) plus the
                      // result of the right subtree, i.e.
                      // dfs(u * 2 + 1, v)
                      return v + this.apply(u * 2, v)
                          + this.apply(u * 2 + 1, v);
                  }
              };

        // We start from node 1 with initial value 0
        return dfs.apply(1, 0);
    }

    public static void main(String[] args)
    {
        int n = 7;
        List<Integer> queries
            = List.of(3, 1, 4, 1, 5, 9, 2, 6);

        int result = numberOfNodes(n, queries);
        System.out.println(result);
    }
}

// This code is contributed by Akshita
Python
def number_of_nodes(n, queries):
    # f[i] = 1 means we need to flip all values in the
    # subtree of the node i
    f = [0] * (n + 1)  # Ensure f has enough space
    # If we flip the node even times, the value would be
    # the same as the original value e.g., 0 (original value) ->
    # 1 -> 0 if we flip the node odd times, the value would
    # be the opposite of the original value e.g., 0
    # (original value) -> 1 -> 0 -> 1. Therefore, we can
    # first process the queries to obtain the final flips
    for q in queries:
        if q <= n:  # Check if query index is within range
            f[q] ^= 1

    def dfs(u, v):
        # u is the current node label
        # If u is greater than n, then return 0
        if u > n:
            return 0
        # Do we need to flip the node u?
        # We flip the value if f[u] is 1
        v ^= f[u]
        # The result would be the value of u, i.e., v
        # plus the result of the left subtree, i.e., dfs(u *
        # 2, v) plus the result of the right subtree, i.e.,
        # dfs(u * 2 + 1, v)
        return v + dfs(u * 2, v) + dfs(u * 2 + 1, v)

    # We start from node 1 with initial value 0
    return dfs(1, 0)

# Driver code
if __name__ == "__main__":
    n = 7
    queries = [3, 1, 4, 1, 5, 9, 2, 6]
    result = number_of_nodes(n, queries)
    print(result)


# This code is contributed by Shivam Gupta
JavaScript
function numberOfNodes(n, queries) {
    // f[i] = 1 means we need to flip all values in the
    // subtree of the node i
    let f = new Array(n + 2).fill(0); // Increase size to n + 2
    // if we flip the node even times, the value would be
    // same as the original value e.g. 0 (original value) ->
    // 1 -> 0 if we flip the node odd times, the value would
    // be the the opposite of the original value e.g.  0
    // (original value) -> 1 -> 0 -> 1 therefore, we can
    // first process the queries to obtain the final flips
    for (let q of queries) {
        f[q] ^= 1;
    }

    function dfs(u, v) {
        // u is the current node label
        // if u is greater than n, then return 0
        if (u > n) {
            return 0;
        }
        // do we need to flip the node u?
        // we flip the value if f[u] is 1
        v ^= f[u];
        // the result would be the value of u, i.e. v
        // plus the result of the left subtree, i.e dfs(u *
        // 2, v) plus the result of the right subtree, i.e.
        // dfs(u * 2 + 1, v)
        return v + dfs(u * 2, v) + dfs(u * 2 + 1, v);
    }

    // we start from node 1 with inital value 0
    return dfs(1, 0);
}

// Driver code
let n = 7;
let queries = [3, 1, 4, 1, 5, 9, 2, 6];
let result = numberOfNodes(n, queries);
console.log(result);

Output
3

Time complexity: O(n), where n is the number of nodes in the tree.
Auxiliary Space: O(n), due to the use of the f array and the recursion stack.


Next Article
Article Tags :
Practice Tags :

Similar Reads