We have the numbers with different colours with the sequence: ['red', 'yellow', 'blue'].

That sequence colours the numbers in the following way:

 1   2   3    4    5   6  7   8   9   10   11  12   13 .....

We have got the following recursive function:

  ```
f(1) = 1
f(n) = f(n - 1) + n
```

Some terms of this sequence with their corresponding colour are:

```
term   value   colour
1        1     "red"
2        3     "blue"
3        6     "blue"
4       10     "red"
5       15     "blue"
6       21     "blue"
7       28     "red"
```

The three terms of the same colour "blue", higher than 3, are: `[6, 15, 21]`

We need a function `same_col_seq(), that may receive three arguments:

- `val`, an integer number
- `k`, an integer number
- `colour`, the name of one of the three colours(red, yellow or blue), as a string.

The function will output a sorted array with the smallest `k` terms, having the same marked colour, but higher than `val`.

Let's see some examples:

```python
same_col_seq(3, 3, 'blue') == [6, 15, 21]
same_col_seq(100, 4, 'red') == [136, 190, 253, 325]
```

The function may output an empty list if it does not find terms of the sequence with the wanted colour in the range [val, 2* k * val]

```python
same_col_seq(250, 6, 'yellow') == []
```

That means that the function did not find any "yellow" term in the range `[250, 3000]`
 
Tests will be with the following features:

* Nmber of tests: `100`
* `100 < val < 1000000`
* `3 < k < 20`

Difficulty: EASY
Tags: Fundamentals

Complete search

Dynamic programming

Mathematics

Data structures
Skill types: Dynamic programming

Data structures

Complete search

Starter code:
def same_col_seq(val, k, col):
	

Solutions:
(D, R) = ({}, [[], [], []])
for i in range(10000):
	D[i] = D.get(i - 1, 0) + i
	R[D[i] % 3].append(D[i])

def same_col_seq(val, k, col):
	r = ['blue', 'red', 'yellow'].index(col)
	return [e for e in R[r] if e > val][:k]


def same_col_seq(val, k, col):
	colDct = {'red': 1, 'blue': 0}

	def gen():
		n = ((1 + 24 * val / 3) ** 0.5 - 1) // 2
		while True:
			n += 1
			s = n * (n + 1) / 2
			if s % 3 == colDct[col]:
				yield s
	g = gen()
	return [next(g) for _ in range(k)] if col != 'yellow' else []


def same_col_seq(val, k, col):
	if col == 'yellow':
		return []
	n = int((2 * val) ** 0.5) - 1
	while n * (n + 1) <= 2 * val:
		n += 1
	values = []
	while len(values) < k:
		if col == 'blue' and (n - 1) % 3 != 0:
			values.append(n * (n + 1) // 2)
		if col == 'red' and (n - 1) % 3 == 0:
			values.append(n * (n + 1) // 2)
		n += 1
	return values


(li, i, pos, d) = (['red', 'yellow', 'blue'], 1, 1, {})
while i < 2000:
	d[pos] = li[(pos - 1) % 3]
	i += 1
	pos += i

def same_col_seq(after, k, need):
	(req, status) = ([], 0)
	for (i, j) in d.items():
		if i > after:
			status = 1
		if status:
			if j == need:
				req.append(i)
		if len(req) == k or i >= 2 * k * after:
			break
	return req


same_col_seq = lambda x, k, c: (lambda n, m: [i * -~i // 2 for i in range(n, n + 3 * k) if i % 3 % 2 == m][:k])(int((8 * x + 1) ** 0.5 + 1) // 2, 'br'.find(c[0]))


def same_col_seq(val, k, col):
	n_start = int(((1 + 8 * val) ** 0.5 - 1) / 2) + 1

	def f(n):
		return (1 + n) * n // 2

	def get_col(n):
		return 'red' if n % 3 == 1 else 'blue'
	return [f(n) for n in range(n_start, n_start + 62) if col == get_col(n)][:k]


from bisect import bisect_right

def same_col_seq(val, k, col, seq=[], colors=[]):
	if not seq:
		color = lambda v: {0: 'blue', 1: 'red', 2: 'yellow'}[v % 3]
		next_val = lambda n: seq[-1] + n
		seq.append(1)
		colors.append('red')
		n = 2
		v = next_val(n)
		while True:
			seq.append(v)
			colors.append(color(v))
			if v > 40 * 10 ** 6:
				break
			n += 1
			v = next_val(n)
	res = []
	index = find_gt(seq, val)
	limit = val * 2 * k
	while True:
		if colors[index] == col:
			res.append(seq[index])
			if len(res) == k:
				break
		if seq[index] > limit and len(res) == 0:
			break
		index += 1
	return res

def find_gt(a, x):
	return bisect_right(a, x)


def same_col_seq(val, k, col):
	results = []
	term = 1
	while values(term) <= val:
		term += 1
	while len(results) < k:
		if colors(values(term)) == col:
			results.append(values(term))
		elif values(term) >= 2 * k * val and len(results) == 0:
			return []
		term += 1
	return results

def values(term):
	if term % 2 == 0:
		return int(term / 2 * term + term / 2)
	else:
		return int((term + 1) / 2 * term)

def colors(term):
	if term % 3 == 0:
		return 'blue'
	else:
		return 'red'


def same_col_seq(val, k, col):
	code = {'blue': 0, 'red': 1, 'yellow': 2}
	result = []
	limit = 2 * k * val
	term = 1
	temp = 1
	while temp <= val:
		temp = temp + term + 1
		term += 1
	while True:
		if temp % 3 == code[col]:
			result.append(temp)
		temp = temp + term + 1
		term += 1
		if len(result) >= k:
			break
		if len(result) == 0 and temp >= limit:
			return []
	return result


import math

def same_col_seq(val, k, col):
	colors = ['red', 'yellow', 'blue']
	m = colors.index(col)
	res = list()
	x0 = math.floor((-1 + math.sqrt(1 + 8 * val)) / 2)
	s = x0 * (1 + x0) // 2
	i = x0 + 1
	while len(res) < k:
		if (s + 2) % 3 == m and s > val:
			res.append(s)
		s += i
		i += 1
		if s > 3 * k * val:
			break
	return res