Median of ancestors for each Node of a given tree

Given a tree with N vertices from 0 to N – 1 (0^th node is root) and val[] where val[i] denotes the value of the i^th vertex, the task is to find the array of integers ans[] of length N, where ans[i] denotes the median value of all its ancestors’ values including ith vertex.

Points to remember:

• If number of nodes are even: then median = ((n/2^th node + ((n)/2^th+1) node) /2
• If the number of nodes is odd: then median = (n+1)/2^th node.

Examples:

Input:

Input 1

val[] = {3, 4, 2, 3, 6, 2, 10, 8, 1}
Output:  3 3.5 2.5 3 4 3 3 3 2 
Explanation:
node 0: {3}, median = 3
node 1: {3, 4}, median = 3.5
node 2: {3, 2}, median = 2.5
node 3: {3, 4, 3}, median = 3
node 4: {3, 4, 6}, mode = 4
node 5: {3, 4, 2}, median = 3
node 6: {3, 2, 10}, median = 3
node 7: {3, 2, 8}, median = 3
node 8: {3, 2, 1}, median = 2

Input:

Input 2

val[] = [1, 3, 2, 1]
Output: 1 2 1.5 1

Approach: The problem can be solved based on the following idea:

Use DFS traversal from root to all nodes in top-down manner and use multisets(s, g) to store the smaller and the greater values from effective median of all the ancestors for a node.

Follow the steps mentioned below to implement the idea:

• Use a top-down DFS to traverse from the root to all the nodes.
• Create multiset s to store values less than the effective median and multiset g for greater values.
• Whenever entering a node, before calling DFS on its children, insert the value of that node in s or g accordingly such that the size of s and g are differed utmost by 1.
• Whenever we exit the node, remove the value of that node from s or g.
• Use a map ans[] to store the required answer for each node.
• If the size of g and s are equal then ans[node] would be the last element of s.
• Otherwise, ans[node] would store the first element of g. 

Below is the implementation of the above approach:

C++

// C++ code to implement the approach

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

// Map to store final ans for each node
unordered_map<int, float> ans;

// MultiSets s and g to store smaller and
// greater values than effective median
multiset<int> s, g;

// Function to add an edge
// between nodes u and v
void addEdge(vector<int> adj[], int u, int v)
{
    adj[u].push_back(v);
    adj[v].push_back(u);
}

// insert v either in s or g
void helper(int v)
{
    s.insert(v);
    auto ls = (--s.end());
    int temp = *ls;
    s.erase(ls);

    g.insert(temp);
    if (g.size() > s.size()) {
        temp = *(g.begin());
        g.erase(g.begin());
        s.insert(temp);
    }
}

// Median of Ancestors
void MedianOfAncestors(vector<int> adj[], int node,
                       vector<int>& vis, vector<int>& val)
{
    int v = val[node];
    vis[node] = 1;

    // insert v into multisets
    helper(v);

    if (g.size() != s.size())
        ans[node] = *(--s.end());
    else
        ans[node]
            = ((*(--s.end())) * 1.0 + (*(g.begin())) * 1.0)
              / 2;

    for (auto child : adj[node]) {
        if (vis[child] == 0)
            MedianOfAncestors(adj, child, vis, val);
    }

    // remove v from multisets
    if (s.find(v) != s.end())
        s.erase(s.find(v));
    else
        g.erase(g.find(v));
}

// Driver Code
int main()
{
    // Number of nodes in a tree
    int N = 9;

    // Initialize tree
    vector<int> adj[N];

    // Tree Formation
    addEdge(adj, 0, 1);
    addEdge(adj, 0, 2);
    addEdge(adj, 1, 3);
    addEdge(adj, 1, 4);
    addEdge(adj, 1, 5);
    addEdge(adj, 2, 6);
    addEdge(adj, 2, 7);
    addEdge(adj, 2, 8);

    // Values of nodes of tree
    vector<int> val = { 3, 4, 2, 3, 6, 2, 10, 8, 1 };
    vector<int> vis(N, 0);

    // Function call
    MedianOfAncestors(adj, 0, vis, val);

    for (int i = 0; i < N; i++)
        cout << ans[i] << " ";
    cout << endl;

    return 0;
}

Java

// Java equivalent
import java.util.*;

class Main
{

  // Map to store final ans for each node
  static HashMap<Integer, Float> ans = new HashMap<>();

  // MultiSets s and g to store smaller and
  // greater values than effective median
  static TreeSet<Integer> s = new TreeSet<>();
  static TreeSet<Integer> g = new TreeSet<>();

  // Function to add an edge
  // between nodes u and v
  static void addEdge(ArrayList<ArrayList<Integer>> adj, int u, int v) {
    adj.get(u).add(v);
    adj.get(v).add(u);
  }

  // insert v either in s or g
  static void helper(int v) {
    s.add(v);
    int temp = s.last();
    s.remove(temp);

    g.add(temp);
    if (g.size() > s.size()) {
      temp = g.first();
      g.remove(temp);
      s.add(temp);
    }
  }

  // Median of Ancestors
  static void MedianOfAncestors(ArrayList<ArrayList<Integer>> adj, int node, ArrayList<Integer> vis, ArrayList<Integer> val) {
    int v = val.get(node);
    vis.set(node, 1);

    // insert v into multisets
    helper(v);

    if (g.size() != s.size())
      ans.put(node, (float) s.last());
    else
      ans.put(node, ((float) s.last() + (float) g.first()) / 2);

    for (int child : adj.get(node)) {
      if (vis.get(child) == 0)
        MedianOfAncestors(adj, child, vis, val);
    }

    // remove v from multisets
    if (s.contains(v))
      s.remove(v);
    else
      g.remove(v);
  }

  // Driver Code
  public static void main(String[] args) {
    // Number of nodes in a tree
    int N = 9;

    // Initialize tree
    ArrayList<ArrayList<Integer>> adj = new ArrayList<>();
    for (int i = 0; i < N; i++)
      adj.add(new ArrayList<>());

    // Tree Formation
    addEdge(adj, 0, 1);
    addEdge(adj, 0, 2);
    addEdge(adj, 1, 3);
    addEdge(adj, 1, 4);
    addEdge(adj, 1, 5);
    addEdge(adj, 2, 6);
    addEdge(adj, 2, 7);
    addEdge(adj, 2, 8);

    // Values of nodes of tree
    ArrayList<Integer> val = new ArrayList<>(Arrays.asList(3, 4, 2, 3, 6, 2, 10, 8, 1));
    ArrayList<Integer> vis = new ArrayList<>(Collections.nCopies(N, 0));

    // Function call
    MedianOfAncestors(adj, 0, vis, val);

    for (int i = 0; i < N; i++)
      System.out.print(ans.get(i) + " ");
    System.out.println();
  }
}

Python3

# Python code to implement the approach
from collections import defaultdict
import heapq

# Map to store final ans for each node
ans = defaultdict(float)

# Heaps s and g to store smaller and
# greater values than effective median
s = []
g = []

# Function to add an edge
# between nodes u and v
def addEdge(adj, u, v):
    adj[u].append(v)
    adj[v].append(u)

# insert v either in s or g
def helper(v):
    s.append(-v)
    heapq.heapify(s)

    temp = -heapq.heappop(s)
    heapq.heappush(g, temp)

    if len(g) > len(s):
        temp = heapq.heappop(g)
        heapq.heappush(s, -temp)

# Median of Ancestors
def MedianOfAncestors(adj, node, vis, val):
    v = val[node]
    vis[node] = 1

    # insert v into heaps
    helper(v)

    if len(g) != len(s):
        ans[node] = -s[0]
    else:
        ans[node] = (g[0] + -s[0]) / 2

    for child in adj[node]:
        if not vis[child]:
            MedianOfAncestors(adj, child, vis, val)

    # remove v from heaps
    if -v in s:
        s.remove(-v)
        heapq.heapify(s)
    else:
        g.remove(v)
        heapq.heapify(g)

# Driver Code


# Number of nodes in a tree
N = 9

# Initialize tree
adj = defaultdict(list)

# Tree Formation
addEdge(adj, 0, 1)
addEdge(adj, 0, 2)
addEdge(adj, 1, 3)
addEdge(adj, 1, 4)
addEdge(adj, 1, 5)
addEdge(adj, 2, 6)
addEdge(adj, 2, 7)
addEdge(adj, 2, 8)

# Values of nodes of tree
val = [3, 4, 2, 3, 6, 2, 10, 8, 1]
vis = [0] * N

# Function call
MedianOfAncestors(adj, 0, vis, val)

for i in range(N):
    print(ans[i], end=" ")
print()

# This Code is contributed by Prasad Kandekar(prasad264)

Javascript

// Importing packages is not necessary in JavaScript
// Map to store final ans for each node
const ans = new Map();

// MultiSets s and g to store smaller and greater values than effective median
const s = new Set();
const g = new Set();

// Function to add an edge between nodes u and v
function addEdge(adj, u, v) {
    adj[u].push(v);
    adj[v].push(u);
    }

// insert v either in s or g
function helper(v) {
    s.add(v);
    let temp = [...s].pop();
    s.delete(temp);

    g.add(temp);
    if (g.size > s.size) {
        temp = [...g][0];
        g.delete(temp);
        s.add(temp);
        }
    }

// Median of Ancestors
function MedianOfAncestors(adj, node, vis, val) {
    let v = val[node];
    vis[node] = 1;

    // insert v into multisets
    helper(v);

    if (g.size !== s.size)
        ans.set(node, s.values().next().value);
    else
        ans.set(node, (s.values().next().value + [...g][0]) / 2);

    for (let child of adj[node]) {
        if (vis[child] === 0)
    MedianOfAncestors(adj, child, vis, val);
}

// remove v from multisets
if (s.has(v))
    s.delete(v);
else
    g.delete(v);
}

// Driver Code
// Number of nodes in a tree
const N = 9;

// Initialize tree
const adj = new Array(N).fill().map(() => []);

// Tree Formation
addEdge(adj, 0, 1);
addEdge(adj, 0, 2);
addEdge(adj, 1, 3);
addEdge(adj, 1, 4);
addEdge(adj, 1, 5);
addEdge(adj, 2, 6);
addEdge(adj, 2, 7);
addEdge(adj, 2, 8);

// Values of nodes of tree
const val = [3, 4, 2, 3, 6, 2, 10, 8, 1];
const vis = new Array(N).fill(0);

// Function call
MedianOfAncestors(adj, 0, vis, val);

for (let i = 0; i < N; i++)
console.log(ans.get(i) + " ");

C#

using System;
using System.Collections.Generic;
using System.Linq;

class MedianOfAncestors {
    static Dictionary<int, double> ans = new Dictionary<
        int, double>(); // Dictionary to store the final
                        // answer for each node
    static List<double> s
        = new List<double>(); // List to store values
                              // smaller than the effective
                              // median
    static List<double> g
        = new List<double>(); // List to store values
                              // greater than the effective
                              // median

    static void addEdge(Dictionary<int, List<int> > adj,
                        int u, int v)
    {
        adj[u].Add(v);
        adj[v].Add(u);
    }

    // Insert v either in s or g
    static void Helper(double v)
    {
        s.Add(-v); // Add the negation of v to s
        s.Sort(); // Sort s in non-increasing order

        double temp = -s[0]; // Get the largest value in s
        s.RemoveAt(0); // Remove the largest value from s
        g.Add(temp); // Add the largest value to g
        g.Sort(); // Sort g in non-decreasing order

        if (g.Count
            > s.Count) // If g has more elements than s,
                       // move an element from g to s
        {
            temp = g[0]; // Get the smallest value in g
            g.RemoveAt(
                0); // Remove the smallest value from g
            s.Add(-temp); // Add the negation of the
                          // smallest value to s
            s.Sort(); // Sort s in non-increasing order
        }
    }

    // Depth First Search function to traverse the tree and
    // compute the effective median of ancestors for each
    // node
    static void DFS(Dictionary<int, List<int> > adj,
                    int node, bool[] vis, int[] val)
    {
        double v = val[node]; // Value of the current node
        vis[node]
            = true; // Mark the current node as visited

        Helper(v); // Insert v either in s or g

        if (g.Count
            != s.Count) // If g has more elements than s,
                        // the effective median is the
                        // largest element in s
        {
            ans[node] = -s[0]; // Store the effective median
                               // for the current node
        }
        else // Otherwise, the effective median is the
             // average of the smallest element in g and the
             // largest element in s
        {
            ans[node] = (g[0] + -s[0])
                        / 2; // Store the effective median
                             // for the current node
        }

        foreach(
            int child in adj[node]) // Traverse the children
                                    // of the current node
        {
            if (!vis[child]) // If the child has not been
                             // visited yet, visit it
            {
                DFS(adj, child, vis, val);
            }
        }

        if (s.Contains(-v)) // If the negation of v is in s,
                            // remove it from s
        {
            s.Remove(-v);
        }
        else // Otherwise, remove v from g
        {
            g.Remove(v);
        }
    }

    static void Main(string[] args)
    {
        int N = 9; // Number of nodes in the tree

        Dictionary<int, List<int> > adj = new Dictionary<
            int,
            List<int> >(); // Dictionary to store the
                           // adjacency list of the tree
        for (int i = 0; i < N; i++) {
            adj[i] = new List<int>();
        }

        addEdge(adj, 0, 1); // Add edges to the tree
        addEdge(adj, 0, 2);
        addEdge(adj, 1, 3);
        addEdge(adj, 1, 4);
        addEdge(adj, 1, 5);
        addEdge(adj, 2, 6);
        addEdge(adj, 2, 7);
        addEdge(adj, 2, 8);

        int[] val
            = new int[] { 3, 4, 2, 3, 6, 2, 10, 8, 1 };
        bool[] vis = new bool[N];

        DFS(adj, 0, vis, val);

        for (int i = 0; i < N; i++) {
            Console.Write(ans[i] + " ");
        }
        Console.WriteLine();
    }
}
// This Code is contributed by Vikram_Shirsat

Output

3 3.5 2.5 3 4 3 3 3 2 

Time Complexity: O(N * log(N))
Auxiliary Space: O(N)