Question
A number of students are taking exams in a room. Students sitting adjacent to each other and taking the same exam can cheat. Arrange the students so that cheating opportunities are minimized. I was free to choose input format.

I chose the input to be a list of length n, denoting n students. The element at index i would indicate the exam student i is taking.
For example, [1,2,3,1,2,2] would mean:

• Student 0 is taking exam 1
• Student 1 is taking exam 2
• Student 2 is taking exam 3
• Student 3 is taking exam 1
• Student 4 is taking exam 2
• Student 5 is taking exam 2

Output would be a list with the students re-arranged. An acceptable output for the above case would be [1,2,3,2,1,2].

Proposed solution
I reduced it to a 2 coloring problem. Adjacent nodes should not have the same color.
I proposed the following algorithm, however I am not sure if it is correct or not. I have so far been unable to find the same question on leetcode.

• Gather the frequencies of all exams
• Put them all in a max heap
• Start with an empty output array
• Pop the highest frequency exam at the beginning
• In a loop:
  • Append to the output array
  • Pop the next highest frequency exam
  • Push the previous exam back with frequency reduced by 1. Do not push if frequency becomes 0

Does this approach work? I tried to test on a few cases and seems like it is working. The interviewer hardly said a word throughout the whole interview, and so I had no idea if I was on the right track.
Here is the code I wrote:

from heapq import heappush,heappop

def arrange(exams):
  n = len(exams)
  if n == 0 or n == 1 or n == 2:
    return exams
  freqs = {}
  for c in exams:
    if c not in freqs:
      freqs[c] = 1
    else:
      freqs[c] += 1
  h = []
  for k,v in freqs.items():
    heappush(h,(-v,k))
  ans = []
  i = 0
  v,k = heappop(h)
  while i < n:
    ans.append(k)
    i += 1
    if h:
      x,y = heappop(h)
    else:
      x,y = v+1,k
    if v+1 < 0:
      heappush(h,(v+1,k))
    v,k = x,y
  return ans