Program find number of subsets of colorful vertices that meets given conditions in Python

Suppose we have an array colors, representing colors for one regular n-gon. Here each vertex of this n-gon was randomly colored with one of n different colors present in the given array. We have to find the number of special subsets of polygon vertices, such that these subsets meet these conditions −

• The size of subset must be at least two.
• If we remove vertices which are present in the subset, from the polygon (the adjacent edges of those vertices will also get removed), then the remaining vertices and edges form some continuous paths.
• None of those paths should contain two vertices of the same color.

We have to count number of such subsets are present. If the answer is too large, then return result mod 10^9 + 7.

So, if the input is like colors = [1,2,3,4], then the output will be 11.

To solve this, we will follow these steps −

• count := an empty map where all values will be an empty list.
• n := size of colors
• for i in range 0 to size of colors - 1, do
  • insert i at the end of count[colors[i]]
• answer := 0
• for i in range 2 to n, do
  • answer := answer + nCr(n, i)
• for each i in list of all keys of count, do
  • l0 := count[i]
  • n0 := size of l0
  • if n0 > 1, then
    • for i in range 0 to n0-2, do
      • for j in range i+1 to n0 - 1, do
        • d1 := l0[j] -l0[i]
        • d2 := l0[i] -l0[j] + n
        • if d1 <= n-3 or d2<= n-3, then
          • answer := answer - 1
• return answer

Example

Let us see the following implementation to get better understanding −

from collections import defaultdict
from math import factorial

def nCr(n, i):
   if n==1:
      return 1
   return factorial(n)//factorial(i)//factorial(n-i)

def solve(colors):
   count = defaultdict(list)
   n = len(colors)

   for i in range(len(colors)):
      count[colors[i]].append(i)
   answer = 0

   for i in range(2, n+1):
      answer += nCr(n, i)

   for i in count.keys():
      l0 = count[i]
      n0 = len(l0)

      if n0 > 1:
         for i in range(n0-1):
            for j in range(i+1, n0):
               d1 = l0[j] -l0[i]
               d2 = l0[i] -l0[j] + n
               if d1 <= n-3 or d2<= n-3:
                  answer -=1

   return answer

colors = [1,2,3,4]
print(solve(colors))

Input

[1,2,3,4]

Output

11