Upper triangular matrix

Function to get the index of a matrix that is storing the upper triangular part of a square matrix.
I found this function in the bottom source, but I had to add the offset.

source: original function without offset

Function

In [1]:

def upper_triangular_index(n, r, c, k=0):
    """
    Returns the index of an array that is storing an
    upper triangular matrix. In this case the matrix
    has to be square and only accepts zero or possitive
    offsets.
    n = square matrix length
    r = actual row
    c = actual column
    k = diagonal possitive offset
    """
    return (n*r-k)+c-((r*(r+1))/2)-r*k

Some examples

In [2]:

import numpy as np

In [3]:

N = 3
keys = range(N)
matrix = np.ones((N,N), dtype=int)*-1

Small example without offset

In [4]:

offset=0
for key1 in keys:
    for key2 in keys:
        if key1+offset <= key2:
            matrix[key1,key2] = \
                upper_triangular_index(N, key1, 
                                       key2, k=offset)
print matrix

[[ 0  1  2]
 [-1  3  4]
 [-1 -1  5]]

Small example with offset = 1

In [5]:

matrix = np.ones((N,N), dtype=int)*-1
offset=1
for key1 in keys:
    for key2 in keys:
        if key1+offset <= key2:
            matrix[key1,key2] = \
                upper_triangular_index(N, key1, 
                                       key2, k=offset)
print matrix

[[-1  0  1]
 [-1 -1  2]
 [-1 -1 -1]]

Large example with offset = 3

In [6]:

N = 9
keys = range(N)
matrix = np.ones((N,N), dtype=int)*-1

In [7]:

offset=3
for key1 in keys:
    for key2 in keys:
        if key1+offset <= key2:
            matrix[key1,key2] = \
              upper_triangular_index(N, key1, 
                                     key2, k=offset)
print matrix

[[-1 -1 -1  0  1  2  3  4  5]
 [-1 -1 -1 -1  6  7  8  9 10]
 [-1 -1 -1 -1 -1 11 12 13 14]
 [-1 -1 -1 -1 -1 -1 15 16 17]
 [-1 -1 -1 -1 -1 -1 -1 18 19]
 [-1 -1 -1 -1 -1 -1 -1 -1 20]
 [-1 -1 -1 -1 -1 -1 -1 -1 -1]
 [-1 -1 -1 -1 -1 -1 -1 -1 -1]
 [-1 -1 -1 -1 -1 -1 -1 -1 -1]]