Graph Coloring (Codeforces 1354E)
Tags: DP, dfs and similar

You are given an undirected graph without self-loops or multiple edges which consists of $$$n$$$ vertices and $$$m$$$ edges. Also you are given three integers $$$n_1$$$, $$$n_2$$$ and $$$n_3$$$.Can you label each vertex with one of three numbers 1, 2 or 3 in such way, that:   Each vertex should be labeled by exactly one number 1, 2 or 3;  The total number of vertices with label 1 should be equal to $$$n_1$$$;  The total number of vertices with label 2 should be equal to $$$n_2$$$;  The total number of vertices with label 3 should be equal to $$$n_3$$$;  $$$|col_u - col_v| = 1$$$ for each edge $$$(u, v)$$$, where $$$col_x$$$ is the label of vertex $$$x$$$. If there are multiple valid labelings, print any of them.

Input: The first line contains two integers $$$n$$$ and $$$m$$$ ($$$1 \le n \le 5000$$$; $$$0 \le m \le 10^5$$$) — the number of vertices and edges in the graph.The second line contains three integers $$$n_1$$$, $$$n_2$$$ and $$$n_3$$$ ($$$0 \le n_1, n_2, n_3 \le n$$$) — the number of labels 1, 2 and 3, respectively. It's guaranteed that $$$n_1 + n_2 + n_3 = n$$$.Next $$$m$$$ lines contan description of edges: the $$$i$$$-th line contains two integers $$$u_i$$$, $$$v_i$$$ ($$$1 \le u_i, v_i \le n$$$; $$$u_i \neq v_i$$$) — the vertices the $$$i$$$-th edge connects. It's guaranteed that the graph doesn't contain self-loops or multiple edges.

Output: If valid labeling exists then print "YES" (without quotes) in the first line. In the second line print string of length $$$n$$$ consisting of 1, 2 and 3. The $$$i$$$-th letter should be equal to the label of the $$$i$$$-th vertex.If there is no valid labeling, print "NO" (without quotes).


Sample Inputs:
6 3
2 2 2
3 1
5 4
2 5

5 9
0 2 3
1 2
1 3
1 5
2 3
2 4
2 5
3 4
3 5
4 5

Sample Outputs:
YES
112323

NO

Solution (Python):
import sys
sys.setrecursionlimit(10 ** 6)
readline = sys.stdin.readline
read = sys.stdin.read
ns = lambda : readline().rstrip()
ni = lambda : int(readline().rstrip())
nm = lambda : map(int, readline().split())
nl = lambda : list(map(int, readline().split()))
prn = lambda x: print(*x, sep='\n')

def solve():
    (n, m) = nm()
    a = nl()
    G = [list() for _ in range(n)]
    for _ in range(m):
        (u, v) = nm()
        u -= 1
        v -= 1
        G[u].append(v)
        G[v].append(u)
    color = [-1] * n
    roots = list()

    def dfs(v, p):
        (siz, bla) = (1, color[v])
        for x in G[v]:
            if x == p:
                continue
            if color[x] < 0:
                color[x] = color[v] ^ 1
                (rs, rb) = dfs(x, v)
                siz += rs
                bla += rb
            elif color[x] == color[v]:
                print('NO')
                exit()
        return (siz, bla)
    for v in range(n):
        if color[v] < 0:
            color[v] = 0
            (siz, bla) = dfs(v, -1)
            roots.append((v, siz, bla))
    dp = [1] * (len(roots) + 1)
    for (i, (v, siz, bla)) in enumerate(roots):
        dp[i + 1] = dp[i] << bla | dp[i] << siz - bla
    if dp[len(roots)] >> a[1] & 1 == 0:
        print('NO')
        return
    cur = a[1]
    ans = [0] * n

    def dfs2(v, p, c):
        if c ^ color[v]:
            if a[0]:
                ans[v] = 1
                a[0] -= 1
            else:
                ans[v] = 3
                a[2] -= 1
        else:
            ans[v] = 2
            a[1] -= 1
        for x in G[v]:
            if not ans[x]:
                dfs2(x, v, c)
    for i in range(len(roots) - 1, -1, -1):
        (v, siz, bla) = roots[i]
        if cur >= bla and dp[i] >> cur - bla & 1:
            fl = 1
            cur -= bla
        else:
            fl = 0
            cur -= siz - bla
        dfs2(v, -1, fl)
    print('YES')
    print(*ans, sep='')
    return
solve()

Solution (Python):
from collections import deque
import sys
input = sys.stdin.readline

def bfs(s, c):
    q = deque()
    q.append(s)
    color[s] = c
    while q:
        i = q.popleft()
        ci = color[i]
        for j in G[i]:
            if not color[j]:
                color[j] = c + 1 if ci == c else c
                q.append(j)
            elif ci == color[j]:
                ans = 'NO'
                print(ans)
                exit()
    return
(n, m) = map(int, input().split())
(n1, n2, n3) = map(int, input().split())
G = [[] for _ in range(n + 1)]
for _ in range(m):
    (u, v) = map(int, input().split())
    G[u].append(v)
    G[v].append(u)
color = [0] * (n + 1)
c = 1
for i in range(1, n + 1):
    if not color[i]:
        bfs(i, c)
        c += 2
l = c // 2
cnt1 = [0] * l
cnt2 = [0] * l
for i in range(1, n + 1):
    ci = color[i]
    if ci % 2:
        cnt1[ci // 2] += 1
    else:
        cnt2[ci // 2 - 1] += 1
dp = [[1]]
for i in range(l):
    (c1, c2) = (cnt1[i], cnt2[i])
    dpi = dp[i]
    l0 = len(dpi) + max(c1, c2)
    dp0 = [0] * l0
    for j in range(len(dpi)):
        dp0[j + c1] |= dpi[j]
        dp0[j + c2] |= dpi[j]
    dp.append(dp0)
ans = 'YES' if len(dp[l]) > n2 and dp[l][n2] else 'NO'
print(ans)
if ans == 'NO':
    exit()
(x, y) = (l, n2)
u = set()
for _ in range(l):
    (c1, c2) = (cnt1[x - 1], cnt2[x - 1])
    if 0 <= y - c1 < len(dp[x - 1]) and dp[x - 1][y - c1]:
        u.add(2 * x - 1)
        y -= c1
    else:
        u.add(2 * x)
        y -= c2
    x -= 1
ans = []
for i in range(1, n + 1):
    if color[i] in u:
        ans.append('2')
    elif n1:
        ans.append('1')
        n1 -= 1
    else:
        ans.append('3')
print(''.join(ans))

Solution (Python):
import sys, math
import io, os
from collections import defaultdict as dd, deque, Counter

def data():
    return sys.stdin.buffer.readline().strip()

def mdata():
    return list(map(int, data().split()))

def outl(var):
    sys.stdout.write(''.join(map(str, var)) + '\n')

def out(var):
    sys.stdout.write(str(var) + '\n')
mod = int(1000000000.0) + 7
(n, m) = mdata()
(n1, n2, n3) = mdata()
g = [[] for i in range(n + 1)]
for i in range(m):
    (u, v) = mdata()
    g[u].append(v)
    g[v].append(u)
vis = [0] * (n + 1)
v = [0] * (n + 1)
ans = [0] * (n + 1)
l1 = [1] * n1 + [3] * n3
dp = []
for i in range(1, n + 1):
    if vis[i] == 0:
        q = [i]
        l = [[], []]
        while q:
            s = q.pop()
            if vis[s] == 1:
                continue
            vis[s] = 1
            l[v[s]].append(s)
            for j in g[s]:
                if vis[j] == 0:
                    v[j] = 1 - v[s]
                    q.append(j)
                elif v[s] - v[j] == 0:
                    out('NO')
                    exit()
        if len(l[1]) > len(l[0]):
            (l[0], l[1]) = (l[1], l[0])
        dp.append(l)
dp.sort(key=lambda x: len(x[0]) - len(x[1]), reverse=True)
for l in dp:
    if max(len(l1), n2) < len(l[0]) or min(len(l1), n2) < len(l[1]):
        out('NO')
        exit()
    if len(l1) > n2:
        for i in l[0]:
            ans[i] = l1.pop()
        for i in l[1]:
            ans[i] = 2
        n2 -= len(l[1])
    else:
        for i in l[1]:
            ans[i] = l1.pop()
        for i in l[0]:
            ans[i] = 2
        n2 -= max(len(l[0]), len(l[1]))
out('YES')
outl(ans[1:])