When started learning programming, I randomly saw this problem on Facebook which now I forgot where the post is. But the problem goes like this:

Note that the order of triplets matter. For example, when \( k=s=2\), the function should return \(6\) since there are \(6\) triplets satisfying the condition: \((2,0,0), (0,2,0), (0,0,2), (1,1,0), (1,0,1) \;\text{and} \; (0,1,1)\).

This is not a tricky question, but I found it interesting for two reasons: first, the statement very simple. Second, it’s a good example for beginners learn how to improve performance. The naive approach is looping through all possible values of \(a, b\) to get an \(O(n^2)\) solution:

def f(k,s):
    count = 0
    for a in range (k+1):
        for b in range(k+1):
            if 0<= s- a - b <=k:
                count += 1       
    return count

I then discussed this problem with my friend, and after several days he gave me a brilliant \(O(n)\) solution:

#O(n) solution
def g(k,s):
    M=min(s,k)
    res=0
    for a in range(k+1):
        if s-a>2*M or s-a<0:
            continue
        else:
            if s-a>=k:
                res+=2*M-(s-a)+1
            else:
                res+=s-a+1
    return res

Let’s me explain his idea: for each \( a\), we count how many pairs \((b,c)\) such that \( 0 \leq b,c \leq k \) and \( b+c=s-a\).

• The main point: we can count pairs \((b,c)\) by the values of \( s-a\):
• As \( s-a \leq k\): the number of pairs increases from \(1\) to \((k+1)\). After that (i.e. when \( s-a = k+1\) it starts decreasing.

This solution is beautiful, isn’t it? However, one should not stop here since there are better solutions.

O(1) solution

It’s easy to see that:

\[ m:=max(0,s-2k)\leq a,b,c \leq M:=min(s,k)\]

Let’s change the point of view: for each fixed \( a \in [m,M]\), and let \( p=b+c=s-a\). For each \( p\), thinking of \( b+c=p\) as a line in the coordinate plane. We want to count how many integer points \( (b,c)\) on the line \( b+c=p\), then sum them up for all \( p\). This can be done by using 1 loop through \( a\)

However, if we know lower bound and upper bound for \( p\), we can calculate the sum without any loops. To do this, we need to study the range of \( p\).
Since \( b+c=s-a\) and \( m\leq a,b,c \leq M \), we always have the “global” lower and upper bound for \( p=b+c=\) :

\[ p_{min}=max(2m,s-M) \leq p=b+c \leq p_{max}=min(2M, s-m)\]

The patterns fall into two scenarios:

If \( s-2k \geq 0\), we have \(M \leq p_{min}\) and \(p_{max}=2M\)

If \( s-2k < 0\), we have \(p_{min}<M\) and \(p_{max} \geq M\)

To see the difference between two cases clearer, we look at the examples in the pictures below.

When \( s-2k < 0\), the green line \(b+c=7\) lies above the diagonal and orange line \(b+c=3\) lies below the diagonal. To count the points inside the the region, we can subtract the total number of points in the upper and lower triangle from the numbers of points in the blue square. In case \( s-2k \geq 0\), we have an arithmetic sum \(1+2+3+\dots\) and this can be computed in constant time. To sum up, let me state the final formula:

The number of triplets \( (a,b,c)\) such that \( 0 \leq a,b,c \leq k \; \text{and} \; a+b+c=s\) is given by \[1+2+…+(p_{\mathrm{max}}-p_{\mathrm{min}}+1)=\frac{1}{2}(p_{\mathrm{max}}-p_{\mathrm{min}}+2)(p_{\mathrm{max}}-p_{\mathrm{min}}+1)\] if \( s-2k \geq 0\), and

\[(M+1)^2- \frac{1}{2}p_{\mathrm{min}}(p_{\mathrm{min}}+1) – \frac{1}{2}(2M-p_{\mathrm{max}})(2M-p_{\mathrm{max}}+1)\]

if \( s-2k < 0\).

Interesting readers may find that the counting integer point problem is somehow related to Pick’s theorem.

Can you think of another \(O(1)\) solution? Yes, there is, and it relies on combinatorics. I will leave it as a fun exercise for math lovers.