Coding Interview Question: Nested List Weight Sum

October 15, 2022

Amazon Interview  Coding & Algorithm Interview  Data Structures & Algorithms  Google Interview  Meta Interview 

Question

Imagine an array that contains both integers and nested arrays, such as the following: [[2, [9, 1, 3], 8], 1, 4]. The depth sum is described as the weighted sum of each integer, weighted by their respective depths. In the example, 9’s depth is 3, while 8’s depth is 2, and 4’s depth is 1.

Given such an array, calculate its depth sum.

For example:

Input = [[2, [9, 1, 3], 8], 1, 4]

Output = 2 * 2 + 3 * (9 + 1 + 3) + 2 * 8 + 1 + 4 = 64

Input = [[[3, [2, 4]]], 3]

Output = 3 * 3 + 4 * (2 + 4) + 3 = 36

Visualization of the first example

C++ Solution

The code contains some extra blocks related to creating random numbers, creating a random nested array, and printing nested arrays.

// https://serhatgiydiren.com
// Oct 15, 2022

#include <bits/stdc++.h>

using namespace std;
using namespace chrono;

mt19937 rng(steady_clock::now().time_since_epoch().count());

struct node
{
  int val;
  vector < node* > children;
  node() : val(0) {}
  node(const int &_val) : val(_val) {}
};

void create_nested_array(node* head, int current_level = 0, int max_level = 0)
{
  if (current_level < max_level)
  {
    if (current_level == 0 || rng() % 3 == 0)
    {
      head->children.resize(1 + (rng() % 3));
      for(auto &e:head->children)
      {
        e = new node();
        create_nested_array(e, current_level + 1, max_level);
      }
    }
    else head->val = 1 + (rng() % 9);
  }
  else head->val = 1 + (rng() % 9);
}

void print(node* head, int cur, int max, int depth)
{
  if (head == NULL) return;
  if (head->children.size()) cout << "[";
  else
  {
    cout << head->val;
    if (cur < max - 1) cout << ", ";
  }
  for(unsigned idx = 0; idx < head->children.size(); idx++) print(head->children[idx], idx, head->children.size(), depth + 1);
  if (head->children.size())
  {
    cout << "]";
    if (cur < max - 1 && depth) cout << ", ";
  }
  if (depth == 0) cout << endl;
}

int calc(node* head, const int &depth = 0)
{
  if (head == NULL) return 0;
  int res = head->val * depth;
  for(auto e:head->children) res += calc(e, depth + 1);
  return res; 
}

void process()
{
  node* root = new node();
  create_nested_array(root, 0, 5);
  print(root, 0, root->children.size(), 0);
  cout << calc(root) << endl;
}

void fast_io()
{
  ios::sync_with_stdio(false);
  srand(time(NULL));
  cin.tie(0);
  cout.tie(0);
  cout << fixed << setprecision(2);
  cerr << fixed << setprecision(2);
}

int main()
{
  fast_io();
  process();
  return EXIT_SUCCESS;
}

Python Solution

// https://serhatgiydiren.com
// Oct 15, 2022

def calc(nested, depth = 0):
  res = 0
  if type(nested) is list:
    for item in nested:
      res += calc(item, depth + 1)
  else:
    res += (nested * depth)
  return res

print(calc([[2, [9, 1, 3], 8], 1, 4])) # 64
print(calc([[[3, [2, 4]]], 3]))        # 36