USACO December 2021 Bronze Problem 1: Lonely Photo (version 2)

# Problem1.py

N = int(input())
lineup = input()
answer = 0

for i in range(N):

    lhs = 0
    if (i>0 and lineup[i-1]!=lineup[i]):
        lhs+=1
        for k in range(i-2,-1,-1):
            if (lineup[k]==lineup[i-1]):
                lhs+=1
            else:
                break

    rhs=0
    if (i+1<N and lineup[i+1]!=lineup[i]):
        rhs+=1
        for k in range(i+2,N):
            if (lineup[k]==lineup[i+1]):
                rhs+=1
            else:
                break

    answer += lhs*rhs + max(lhs-1,0) + max(rhs-1,0)
                
print(answer, end="")

USACO December 2021 Bronze Problem 1: Lonely Photo (version 1)

/*  prob1.java
*
*/

import java.io.*;
import java.util.*;

class prob1 {
    public static void main (String [] args) throws IOException {

        BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
        PrintWriter out = new PrintWriter(System.out);
        StringTokenizer st;
    
        st = new StringTokenizer(in.readLine());         
        int N = Integer.parseInt(st.nextToken());        
        // out.println(N);  // DEBUG
        
        int[] arrH = new int[N];                       // Holstein zero-one array
        int[] arrHsum = new int[N];                    // Holstein subset sum array

        st = new StringTokenizer(in.readLine());         
        char[] arrinput = (st.nextToken()).toCharArray();
         
        long answer = 0;
        int i,j;
        char c;
        int seqlen;
        int hcount;
        
        // initialize the subset sum
        if (arrinput[0] == 'H') {
			arrH[0]=1;
			arrHsum[0]=1;
		}
        
        // loop to fill in the subset sum
        for (i = 1; i < N; i++) {
            if(arrinput[i] == 'H') {
				arrH[i] = 1;
			} 
			arrHsum[i] = arrHsum[i-1] + arrH[i];
        }
        
        // loop to check sequences that start at index 0
        for (j=2;j<N;j++) {
			seqlen = j+1;
			if ((arrHsum[j]==1)||(arrHsum[j]==seqlen-1))  {
			    answer++;	
			}
	}
		
	// nested loop to check all other sequences
	for (i=1;i<N-2;i++) {
		for (j=i+2;j<N;j++) {
			seqlen = j-i+1;
		        // sequence starts at index i (for i>0)
		        // sequence ends at index j
			hcount = arrHsum[j]-arrHsum[i-1];
			if ((hcount==1)||(hcount==seqlen-1)) {
			    answer++;	
			}
		}		
	}
        
        // out.println(Arrays.toString(arrinput));          // DEBUG    
        // out.println(Arrays.toString(arrH));              // DEBUG
        // out.println(Arrays.toString(arrHsum));           // DEBUG
        
        out.println(answer);
        out.close();
    }
}