Sum of max elements in sub-triangles

Question

I attempted this last question from this year's CCC 2019 S5.

Problem statement:

In a parallel universe, the most important data structure in computer science is the triangle. A triangle of size M consists of M rows, with the ith row containing i elements. Furthermore, these rows must be arranged to form the shape of an equilateral triangule. That is, each row is centred around a vertical line of symmetry through the middle of the triangle. For example, the diagram below shows a triangle of size 4:

A triangle contains sub-triangles. For example, the triangle above contains ten sub-triangles of size 1, six sub-triangles of size 2 (two of which are the triangle containing (3,1,2) and the triangle containing (4,6,1)), three sub-triangles of size 3 (one of which contains (2,2,1,1,4,2)). Note that every triangle is a sub-triangle of itself.

You are given a triangle of size N and must find the sum of the maximum elements of every sub-triangle of size K.

Input Specification

The first line contains two space-separated integers N and K (1≤K≤N≤3000).

Following this are N lines describing the triangle. The ith of these lines contains i space-separated integers ai,j (0≤ai,j≤10^9), representing the ith row of the triangle.

For 4 of the 15 available marks, N≤1000.

Output Specification

Output the integer sum of the maximum elements of every sub-triangle of size K.

Sample Input

4 2
3
1 2
4 2 1
6 1 4 2

Output for Sample Input

23

Unfortunately my solution gave TLE verdict, and I don't know how to optimize it.

The question basically asks to find maximum elements of subtriangles and add them together. My approach's straightforward, I iterate each elements of the large triangle making them a "root" of a subtriangle then I try to go to each of its elements to find max and add them to the result.

I need some help on improving my solution, does it require some data structure?

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    vector<vector<int>> triangle;

    int n;
    int k;

    cin >> n >> k;

    for (int i = 0; i < n; ++i)
    {
        triangle.push_back(vector<int>(i + 1, 0));

        for (int j = 0; j <= i; ++j)
            cin >> triangle[i][j];
    }

    int sum = 0;

    for (int level = 0; level <= n - k; ++level)
        for (int root = 0, size1 = triangle[level].size(); root < size1; ++root)
        {
            int largest = 0;

            for (int i = level, counter = 0; i < level + k; ++i, ++counter)
                for (int j = root, times = 0, size2 = triangle[i].size(); j < size2 && times <= counter; ++j, ++times)
                    largest = max(largest, triangle[i][j]);

            sum += largest;
        }

    cout << sum << "\n";

    return 0;
}

Answer 1

Here is a solution that can be made O(n^2 log(k)) which is fast enough.

The idea is this. Going from the nxn triangle of triangles of size 1 to the (n-1)x(n-1) triangle of max values of triangles of size 2 is a O(n) operation. Just compare each triangle to the max of its neighbors.

The same trick can be used to go from that second triangle to the (n-2)x(n-2) triangle of triangles of size 2. But instead if you skip one in each direction, you can get directly to the (n-3)x(n-3) triangle of max values in triangles of size 4. Also in time O(n). To illustrate the latter suppose that we started with:

    2
   3 1
  1 2 4
 4 2 1 5
6 1 4 2 3

To get to the size 2 triangles we compare each triangle to its neighbors.

   3
  3 4
 4 2 5
6 4 4 5

And to get to the size 4 triangle compare skipping one, so the bottom one we compare 6, 3, 4. The next over we compare 4, 4, 5 and so on. To get:

 5
6 5

And then we add them together to get 11.

Next, from the (n-3)x(n-3) triangle of max values in triangles of size 4 you can go directly to the triangle of max values in triangles of sizes 5, 6, 7, or 8 by choosing the size of the triangles we'll compare to be next, skip 1, skip 2 or skip 3.

And so on resulting in O(log(k)) steps to get the triangle of max values in k by k triangles.

Answer 2

Here is my solution:

#include <iostream>
#include <vector>

using namespace std;

int max3(int a, int b, int c)
{
    return max(a, max(b,c));
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n,k;
    cin>>n>>k;
    vector<vector<long long>> triangle;
    for(int i = 0; i < n; i++)
    {
        triangle.push_back(vector<long long>(i + 1, 0));
        for (int j = 0; j <= i; j++)
            cin >> triangle[i][j];
    }
    for(int size = 2; size <= k; size++)
        for(int i = 0;i <= n - size; i++)
            for(int j = 0; j <= i; j++)
                triangle[i][j] = max3(triangle[i][j], triangle[i+1][j+1], triangle[i+1][j]);
    long long sum = 0;
    for(int i = 0;i <= n - k; i++)
        for(int j = 0; j <= i; j++)
            sum += triangle[i][j];
    cout<<sum;
} 

Explanation: I will call the upmost square of a triangle the top of the triangle. May t(i, j, k) be the biggest number in a triangle with top at i, j and size k. From this we observe the following facts:

• t(i, j, 1) = triangle[i][j]
• Every triangle with top at (i, j) of size k>=2 can be made by taking the union of 3 other triangles of size k-1 with tops at (i, j), (i+1, j) and (i+1, j+1)
• Therefore the biggest number in a triangle with top at (i, j) of size k>=2 is going to be the maximum of the max numbers of these 3 triangles or written as a formula:
• t(i, j, k) = max( t(i, j, k-1), t(i+1, j, k-1), t(i, j+1, k-1) )

So all we need to do is store the maximum values of the triangle for the previous k. Since we are going to iterate the triangle starting from the top we can overwrite the value currently in it since the formula only uses triangles under it and itself to calculate the value for the new k. The program should start from size = 2 and calculate the new value for all sizes until k by using the old values. I also use long long cause it's better to be safe than sorry. Hope this helps!